diff --git a/pyproject.toml b/pyproject.toml index d7a96faa..573d8faa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -95,7 +95,7 @@ colab_evals = [ # Statistical Analysis # Data Manipulation & Display - "unstructured>=0.18.15", + "unstructured>=0.18.15", ] colab_fit = [ "pandas>=2.0.0,<2.3.0", # Colab compatibility (2.2.2) and cudf constraint @@ -109,9 +109,9 @@ colab_fit = [ "evaluate>=0.4.5", "rouge-score>=0.1.2", "sentencepiece>=0.2.1", + "ray[default]>=2.46.0", # Ray for distributed training "mlflow>=3.2.0", "requests>=2.32.0", # Relaxed for Colab (2.32.4) - "loguru>=0.7.3", "ipython>=7.34.0", # Colab compatibility (7.34.0) "jupyter>=1.1.1", "ipywidgets>=7.3.4,<9.0.0", # Support both v7 (Colab) and v8 (Jupyter) @@ -178,7 +178,6 @@ local_evals = [ "sentencepiece>=0.2.1", "mlflow>=3.2.0", "requests>=2.32.0", # Relaxed for Colab (2.32.4) - "loguru>=0.7.3", "ipython>=7.34.0", # Colab compatibility (7.34.0) "jupyter>=1.1.1", "flashinfer-python==0.2.5", @@ -245,7 +244,6 @@ local_fit = [ "dill>=0.3.8", "mlflow>=3.2.0", "requests>=2.32.0", # Relaxed for Colab (2.32.4) - "loguru>=0.7.3", "ipython>=7.34.0", # Colab compatibility (7.34.0) "jupyter>=1.1.1", "trackio", diff --git a/rapidfireai/automl/automl_utils.py b/rapidfireai/automl/automl_utils.py index 83eba86a..a0b9f1ae 100644 --- a/rapidfireai/automl/automl_utils.py +++ b/rapidfireai/automl/automl_utils.py @@ -3,7 +3,7 @@ from typing import Any from rapidfireai.automl.base import AutoMLAlgorithm -from rapidfireai.fit.utils.exceptions import AutoMLException +from rapidfireai.utils.exceptions import AutoMLException def _is_valid_reranker_top_n_vs_k(pipeline: Any) -> bool: diff --git a/rapidfireai/automl/base.py b/rapidfireai/automl/base.py index e9b18d9d..04421464 100644 --- a/rapidfireai/automl/base.py +++ b/rapidfireai/automl/base.py @@ -4,7 +4,7 @@ from typing import Any from rapidfireai.automl.datatypes import List -from rapidfireai.fit.utils.exceptions import AutoMLException +from rapidfireai.utils.exceptions import AutoMLException class AutoMLAlgorithm(ABC): diff --git a/rapidfireai/automl/grid_search.py b/rapidfireai/automl/grid_search.py index 7f9090b6..16bcabe4 100644 --- a/rapidfireai/automl/grid_search.py +++ b/rapidfireai/automl/grid_search.py @@ -7,7 +7,7 @@ from rapidfireai.automl.base import AutoMLAlgorithm from rapidfireai.automl.datatypes import List from rapidfireai.automl.automl_utils import filter_evals_runs_valid_reranker -from rapidfireai.fit.utils.exceptions import AutoMLException +from rapidfireai.utils.exceptions import AutoMLException def recursive_expand_gridsearch(item: Any): diff --git a/rapidfireai/automl/model_config.py b/rapidfireai/automl/model_config.py index 775596f5..9defcedd 100644 --- a/rapidfireai/automl/model_config.py +++ b/rapidfireai/automl/model_config.py @@ -6,7 +6,7 @@ from abc import ABC, abstractmethod from collections.abc import Callable from dataclasses import dataclass -from typing import Any, Type, get_type_hints +from typing import Any, get_type_hints from rapidfireai.automl.datatypes import List, Range @@ -207,7 +207,7 @@ def get_engine_kwargs(self) -> dict[str, Any]: pass -def _create_rf_class_evals(base_class: Type, class_name: str): +def _create_rf_class_evals(base_class: type, class_name: str): """Creating a RF class for evals that dynamically inherits all constructor parameters and supports singleton, list, and Range values.""" if not inspect.isclass(base_class): raise ValueError(f"base_class must be a class, got {type(base_class)}") diff --git a/rapidfireai/automl/random_search.py b/rapidfireai/automl/random_search.py index 171bfd51..b2caede1 100644 --- a/rapidfireai/automl/random_search.py +++ b/rapidfireai/automl/random_search.py @@ -1,14 +1,14 @@ """Random search implementation for AutoML hyperparameter optimization.""" +import hashlib import json import random -import hashlib from typing import Any from rapidfireai.automl.base import AutoMLAlgorithm from rapidfireai.automl.datatypes import List, Range from rapidfireai.automl.automl_utils import _is_valid_reranker_top_n_vs_k -from rapidfireai.fit.utils.exceptions import AutoMLException +from rapidfireai.utils.exceptions import AutoMLException def encode_payload(payload: dict[str, Any]) -> str: @@ -73,16 +73,12 @@ def _get_runs_fit(self, seed: int) -> list[dict[str, Any]]: selected_peft_config = config.peft_config peft_params = ( - {} - if selected_peft_config is None - else recursive_expand_randomsearch(selected_peft_config._user_params) + {} if selected_peft_config is None else recursive_expand_randomsearch(selected_peft_config._user_params) ) # Sample other parameters training_params = ( - {} - if config.training_args is None - else recursive_expand_randomsearch(config.training_args._user_params) + {} if config.training_args is None else recursive_expand_randomsearch(config.training_args._user_params) ) model_kwargs = ( @@ -202,11 +198,7 @@ def _get_runs_evals(self, seed: int) -> list[dict[str, Any]]: for pipeline in pipelines: # Sample model config parameters - pipeline_instances = ( - [{}] - if pipeline is None - else [recursive_expand_randomsearch(pipeline)] - ) + pipeline_instances = [{}] if pipeline is None else [recursive_expand_randomsearch(pipeline)] additional_kwargs = { k: v @@ -217,9 +209,7 @@ def _get_runs_evals(self, seed: int) -> list[dict[str, Any]]: and v is not None } additional_kwargs_instances = ( - [{}] - if not additional_kwargs - else [recursive_expand_randomsearch(additional_kwargs)] + [{}] if not additional_kwargs else [recursive_expand_randomsearch(additional_kwargs)] ) # Generate random search combinations diff --git a/rapidfireai/cli.py b/rapidfireai/cli.py index 7641b505..055b0d30 100644 --- a/rapidfireai/cli.py +++ b/rapidfireai/cli.py @@ -5,7 +5,6 @@ import argparse import os -import platform import re import signal import shutil @@ -13,13 +12,20 @@ import subprocess import sys from pathlib import Path -from importlib.resources import files -from rapidfireai.utils.get_ip_address import get_ip_address -from rapidfireai.utils.python_info import get_python_info -from rapidfireai.utils.constants import DispatcherConfig, JupyterConfig, ColabConfig -from rapidfireai.utils.doctor import get_doctor_info -from rapidfireai.utils.constants import RF_EXPERIMENT_PATH, RF_HOME -from rapidfireai.utils.gpu_info import get_compute_capability + +from rapidfireai.utils.constants import ( + RF_DB_PATH, + RF_EXPERIMENT_PATH, + RF_HOME, + RF_LOG_PATH, + ColabConfig, + DispatcherConfig, + JupyterConfig, +) +from rapidfireai.platform.doctor import get_doctor_info +from rapidfireai.platform.get_ip_address import get_ip_address +from rapidfireai.platform.gpu_info import get_compute_capability +from rapidfireai.platform.python_info import get_python_info from .version import __version__ @@ -110,41 +116,27 @@ def get_cuda_version(): return 0, 0 -def install_packages(evals: bool = False, init_packages: list[str] | None = None): +def install_packages(init_packages: list[str] | None = None): """Install packages for the RapidFire AI project.""" packages = [] - # Generate CUDA requirements file - mode_file = Path(RF_HOME) / "rf_mode.txt" - if evals: - mode_file.write_text("evals") - else: - mode_file.write_text("fit") cuda_major, cuda_minor = get_cuda_version() python_info = get_python_info() site_packages = python_info["site_packages"] setup_directory = None for site_package in site_packages.split(",") + ["."]: - if os.path.exists(os.path.join(site_package.strip(), "setup", "fit")): + if os.path.exists(os.path.join(site_package.strip(), "setup", "rapidfireai")): setup_directory = Path(site_package) / "setup" break if not setup_directory: print("โŒ Setup directory not found, skipping package installation") return 1 - if ColabConfig.ON_COLAB and evals: - print("Colab environment detected, installing evals packages") - requirements_file = setup_directory / "evals" / "requirements-colab.txt" - elif ColabConfig.ON_COLAB and not evals: - print("Colab environment detected, installing fit packages") - requirements_file = setup_directory / "fit" / "requirements-colab.txt" - elif not ColabConfig.ON_COLAB and evals: - print("Non-Colab environment detected, installing evals packages") - requirements_file = setup_directory / "evals" / "requirements-local.txt" - elif not ColabConfig.ON_COLAB and not evals: - print("Non-Colab environment detected, installing fit packages") - requirements_file = setup_directory / "fit" / "requirements-local.txt" + + if ColabConfig.ON_COLAB: + print("Colab environment detected, installing packages") + requirements_file = setup_directory / "rapidfireai" / "requirements-colab.txt" else: - print("โŒ Unknown environment detected, skipping package installation") - return 1 + print("Local/server environment detected, installing packages") + requirements_file = setup_directory / "rapidfireai" / "requirements-local.txt" try: print(f"Installing packages from {requirements_file.absolute()}...") @@ -167,8 +159,8 @@ def install_packages(evals: bool = False, init_packages: list[str] | None = None torchaudio_version = "2.5.1" torch_cuda = "cu121" flash_cuda = "cu121" - if cuda_major==12: - if cuda_minor>=9: + if cuda_major == 12: + if cuda_minor >= 9: # Supports Torch 2.8.0 torch_version = "2.8.0" torchvision_version = "0.23.0" @@ -177,7 +169,7 @@ def install_packages(evals: bool = False, init_packages: list[str] | None = None flash_cuda = "cu129" vllm_cuda = "cu129" vllm_version = "0.11.0" - elif cuda_minor>=8: + elif cuda_minor >= 8: # Supports Torch 2.9.0/1 torch_version = "2.8.0" torchvision_version = "0.23.0" @@ -186,7 +178,7 @@ def install_packages(evals: bool = False, init_packages: list[str] | None = None flash_cuda = "cu128" vllm_cuda = "cu128" vllm_version = "0.11.0" - elif cuda_minor>=6: + elif cuda_minor >= 6: # Supports Torch 2.9.0/1 torch_version = "2.8.0" torchvision_version = "0.23.0" @@ -194,7 +186,7 @@ def install_packages(evals: bool = False, init_packages: list[str] | None = None torch_cuda = "cu126" flash_cuda = "cu126" vllm_cuda = "cu126" - elif cuda_minor>=4: + elif cuda_minor >= 4: # Supports Torch 2.6.0 torch_version = "2.6.0" torchvision_version = "0.21.0" @@ -212,7 +204,7 @@ def install_packages(evals: bool = False, init_packages: list[str] | None = None flash_cuda = "cu121" vllm_cuda = "cu121" - elif cuda_major==13: + elif cuda_major == 13: # Supports Torch 2.9.0/1 torch_version = "2.8.0" torchvision_version = "0.23.0" @@ -227,40 +219,47 @@ def install_packages(evals: bool = False, init_packages: list[str] | None = None if ColabConfig.ON_COLAB: flash_cuda = "cu128" - if not evals: - pass - - if evals and ColabConfig.ON_COLAB: - pass - - - ## TODO: re-enable for fit once trl has fix if not ColabConfig.ON_COLAB and cuda_major >= 12: print(f"\n๐ŸŽฏ Detected CUDA {cuda_major}.{cuda_minor}, using {torch_cuda}") - + + packages.append( + { + "package": f"torch=={torch_version}", + "extra_args": ["--upgrade", "--index-url", f"https://download.pytorch.org/whl/{torch_cuda}"], + } + ) + packages.append( + { + "package": f"torchvision=={torchvision_version}", + "extra_args": ["--upgrade", "--index-url", f"https://download.pytorch.org/whl/{torch_cuda}"], + } + ) + packages.append( + { + "package": f"torchaudio=={torchaudio_version}", + "extra_args": ["--upgrade", "--index-url", f"https://download.pytorch.org/whl/{torch_cuda}"], + } + ) + packages.append({"package": f"vllm=={vllm_version}", "extra_args": ["--upgrade"]}) + packages.append({"package": "flashinfer-python", "extra_args": ["--upgrade"]}) + packages.append({"package": "flashinfer-cubin", "extra_args": ["--upgrade"]}) + if cuda_major + (cuda_minor / 10.0) >= 12.8: + packages.append( + { + "package": "flashinfer-jit-cache", + "extra_args": ["--upgrade", "--index-url", f"https://flashinfer.ai/whl/{flash_cuda}"], + } + ) + packages.append({"package": "transformers>=4.56.1,<5.0.0", "extra_args": ["--upgrade"]}) packages.append({"package": f"torch=={torch_version}", "extra_args": ["--upgrade", "--index-url", f"https://download.pytorch.org/whl/{torch_cuda}"]}) packages.append({"package": f"torchvision=={torchvision_version}", "extra_args": ["--upgrade", "--index-url", f"https://download.pytorch.org/whl/{torch_cuda}"]}) packages.append({"package": f"torchaudio=={torchaudio_version}", "extra_args": ["--upgrade", "--index-url", f"https://download.pytorch.org/whl/{torch_cuda}"]}) - if evals: - packages.append({"package": f"vllm=={vllm_version}", "extra_args": ["--upgrade"]}) - packages.append({"package": "flashinfer-python", "extra_args": ["--upgrade"]}) - packages.append({"package": "flashinfer-cubin", "extra_args": ["--upgrade"]}) - if cuda_major + (cuda_minor / 10.0) >= 12.8: - packages.append({"package": "flashinfer-jit-cache", "extra_args": ["--upgrade","--index-url", f"https://flashinfer.ai/whl/{flash_cuda}"]}) - # Re-install torch, torchvision, and torchaudio to ensure compatibility as many packages try and upgrade it - packages.append({"package": "transformers>=4.56.1,<5.0.0", "extra_args": ["--upgrade"]}) + if get_compute_capability() >= 8.0: + packages.append({"package": "flash-attn>=2.8.3", "extra_args": ["--upgrade", "--no-build-isolation"]}) + # Re-install torch after flash-attn to ensure compatibility packages.append({"package": f"torch=={torch_version}", "extra_args": ["--upgrade", "--index-url", f"https://download.pytorch.org/whl/{torch_cuda}"]}) packages.append({"package": f"torchvision=={torchvision_version}", "extra_args": ["--upgrade", "--index-url", f"https://download.pytorch.org/whl/{torch_cuda}"]}) packages.append({"package": f"torchaudio=={torchaudio_version}", "extra_args": ["--upgrade", "--index-url", f"https://download.pytorch.org/whl/{torch_cuda}"]}) - if get_compute_capability() >= 8.0: - packages.append({"package": "flash-attn>=2.8.3", "extra_args": ["--upgrade", "--no-build-isolation"]}) - # Re-install torch, torchvision, and torchaudio to ensure compatibility as flash-attn requires an old version of torch but will upgrade torch to an incompatible version - packages.append({"package": f"torch=={torch_version}", "extra_args": ["--upgrade", "--index-url", f"https://download.pytorch.org/whl/{torch_cuda}"]}) - packages.append({"package": f"torchvision=={torchvision_version}", "extra_args": ["--upgrade", "--index-url", f"https://download.pytorch.org/whl/{torch_cuda}"]}) - packages.append({"package": f"torchaudio=={torchaudio_version}", "extra_args": ["--upgrade", "--index-url", f"https://download.pytorch.org/whl/{torch_cuda}"]}) - # else: - # packages.append({"package": "flash-attn-triton", "extra_args": ["--upgrade"]}) - # packages.append({"package": "https://github.com/RapidFireAI/faiss-wheels/releases/download/v1.13.0/rf_faiss_gpu_12_8-1.13.0-cp39-abi3-manylinux_2_34_x86_64.whl", "extra_args": []}) packages.append({"package": "numpy<2.3", "extra_args": ["--upgrade"]}) @@ -301,16 +300,17 @@ def copy_tutorial_notebooks(): return 0 -def run_init(evals: bool = False): +def run_init(): """Run the init command to initialize the project.""" print("๐Ÿ”ง Initializing RapidFire AI project...") print("-" * 30) - print("Initializing project...") - install_packages(evals) + install_packages() copy_tutorial_notebooks() - + print("-" * 30) + print("โœ… RapidFire AI initialization complete!") return 0 + def copy_test_notebooks(): """Copy the test notebooks to the project.""" print("Getting test notebooks...") @@ -329,12 +329,44 @@ def copy_test_notebooks(): return 1 return 0 + +def run_clear(): + """Clear the database and all log files.""" + print("๐Ÿงน Clearing RapidFire AI data...") + print("-" * 30) + + paths_to_clear = [ + ("Database", RF_DB_PATH), + ("Logs", RF_LOG_PATH), + ("Experiments", RF_EXPERIMENT_PATH), + ] + + # Delete existing directories + deleted_any = False + for name, path in paths_to_clear: + if os.path.exists(path): + try: + shutil.rmtree(path) + print(f"โœ… Deleted {name}: {path}") + deleted_any = True + except Exception as e: + print(f"โŒ Failed to delete {name}: {e}") + + if not deleted_any: + print("โœ… Nothing to clear - all directories are already empty or don't exist.") + + print("-" * 30) + print("โœ… RapidFire AI data cleared!") + return 0 + + def run_jupyter(): - """ Run the Jupyter notebook server. """ - from jupyter_server.serverapp import ServerApp - import logging + """Run the Jupyter notebook server.""" import io - from contextlib import redirect_stdout, redirect_stderr + import logging + from contextlib import redirect_stderr, redirect_stdout + + from jupyter_server.serverapp import ServerApp # Suppress all logging logging.getLogger().setLevel(logging.CRITICAL) @@ -344,9 +376,9 @@ def run_jupyter(): app = ServerApp() app.open_browser = False app.port = JupyterConfig.PORT - app.allow_origin = '*' + app.allow_origin = "*" app.websocket_ping_interval = 90000 - app.log_level = 'CRITICAL' + app.log_level = "CRITICAL" app.token = "" app.password = "" app.default_url = "/tree" @@ -356,8 +388,8 @@ def run_jupyter(): try: with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture): - app.initialize(argv=['--ServerApp.custom_display_url=']) - + app.initialize(argv=["--ServerApp.custom_display_url="]) + dispatcher_port = DispatcherConfig.PORT if os.getenv("TERM_PROGRAM") == "vscode": @@ -367,63 +399,75 @@ def run_jupyter(): os_username = os.getenv("USER", os.getenv("LOGNAME", "username")) print(f"Manually forward port {app.port} to localhost") print(f"Manually forward port {dispatcher_port} to localhost") - print(f"For example using ssh:") - print(f" ssh -L {app.port}:localhost:{app.port} -L {dispatcher_port}:localhost:{dispatcher_port} {os_username}@{get_ip_address()}") + print("For example using ssh:") + print( + f" ssh -L {app.port}:localhost:{app.port} -L {dispatcher_port}:localhost:{dispatcher_port} {os_username}@{get_ip_address()}" + ) print("If there is a problem, try running jupyter manually with:") - print(f" jupyter notebook --no-browser --port={app.port} --ServerApp.allow_origin='*' --ServerApp.default_url='/tree' --ServerApp.token=''") + print( + f" jupyter notebook --no-browser --port={app.port} --ServerApp.allow_origin='*' --ServerApp.default_url='/tree' --ServerApp.token=''" + ) print("\n\nAfter forwarding the ports above, access the Jupyter notebook at:") print(f"http://localhost:{app.port}/tree?token={app.token}") - + # Don't redirect anything during start - let prompts through app.start() - + except Exception as e: print("ERROR occurred during Jupyter server startup:", file=sys.stderr) print("=" * 60, file=sys.stderr) - + stdout_output = stdout_capture.getvalue() stderr_output = stderr_capture.getvalue() - + if stdout_output: print(" Standard output:", file=sys.stderr) print(stdout_output, file=sys.stderr) - + if stderr_output: print(" Standard error:", file=sys.stderr) print(stderr_output, file=sys.stderr) - + print("=" * 60, file=sys.stderr) print(f"Exception: {e}", file=sys.stderr) print("Try running jupyter manually with:") - print(f" jupyter notebook --no-browser --port={app.port} --ServerApp.allow_origin='*' --ServerApp.default_url='/tree' --ServerApp.token=''") + print( + f" jupyter notebook --no-browser --port={app.port} --ServerApp.allow_origin='*' --ServerApp.default_url='/tree' --ServerApp.token=''" + ) raise + def main(): """Main entry point for the rapidfireai command.""" - parser = argparse.ArgumentParser(description="RapidFire AI - Start/stop/manage services", prog="rapidfireai", - epilog=""" + parser = argparse.ArgumentParser( + description="RapidFire AI - Start/stop/manage services", + prog="rapidfireai", + epilog=""" Examples: - # Basic initialization for training + # Initialize RapidFire AI (installs all dependencies) rapidfireai init - #or - # Basic Initialize with evaluation dependencies - rapidfireai init --evals - + # Start services rapidfireai start - + # Stop services rapidfireai stop + # Diagnose issues + rapidfireai doctor + + # Clear database and log files + rapidfireai clear + For more information, visit: https://github.com/RapidFireAI/rapidfireai - """ + """, ) parser.add_argument( "command", nargs="?", default="start", - choices=["start", "stop", "status", "restart", "setup", "doctor", "init", "jupyter"], + choices=["start", "stop", "status", "restart", "setup", "doctor", "init", "jupyter", "clear"], help="Command to execute (default: start)", ) @@ -458,9 +502,7 @@ def main(): parser.add_argument("--force", "-f", action="store_true", help="Force action without confirmation") - parser.add_argument("--evals", action="store_true", help="Initialize with evaluation dependencies") - - parser.add_argument("--log-lines", type=int, default=10, help="Number of lines to log to the console") + parser.add_argument("--log-lines", type=int, default=10, help="Number of log lines to show in doctor command") parser.add_argument( "--converge", @@ -485,11 +527,9 @@ def main(): os.environ["RF_TRACKIO_ENABLED"] = "true" if args.tensorboard_log_dir: os.environ["RF_TENSORBOARD_LOG_DIR"] = args.tensorboard_log_dir - if args.colab: - os.environ["RF_COLAB_MODE"] = "true" - elif ColabConfig.ON_COLAB and os.getenv("RF_COLAB_MODE") is None: + if args.colab or (ColabConfig.ON_COLAB and os.getenv("RF_COLAB_MODE") is None): os.environ["RF_COLAB_MODE"] = "true" - + # Handle force command separately if args.force: os.environ["RF_FORCE"] = "true" @@ -503,11 +543,14 @@ def main(): # Handle init command separately if args.command == "init": - return run_init(args.evals) - + return run_init() + if args.command == "jupyter": return run_jupyter() + if args.command == "clear": + return run_clear() + if args.test_notebooks: return copy_test_notebooks() diff --git a/rapidfireai/db/__init__.py b/rapidfireai/db/__init__.py new file mode 100644 index 00000000..173b1202 --- /dev/null +++ b/rapidfireai/db/__init__.py @@ -0,0 +1,10 @@ +""" +RapidFire AI Database Module + +Provides SQLite database interface and operations for experiment tracking. +""" + +from rapidfireai.db.db_interface import DatabaseInterface +from rapidfireai.db.rf_db import RfDb + +__all__ = ["DatabaseInterface", "RfDb"] diff --git a/rapidfireai/evals/db/db_interface.py b/rapidfireai/db/db_interface.py similarity index 70% rename from rapidfireai/evals/db/db_interface.py rename to rapidfireai/db/db_interface.py index 6b5a6090..57708485 100644 --- a/rapidfireai/evals/db/db_interface.py +++ b/rapidfireai/db/db_interface.py @@ -1,4 +1,8 @@ -"""Interface for the database.""" +""" +RapidFire AI Database Interface + +Provides low-level SQLite database connection and query execution. +""" import functools import os @@ -7,27 +11,57 @@ from collections.abc import Callable from typing import Any -from rapidfireai.evals.utils.constants import DBConfig +from rapidfireai.utils.constants import RF_DB_PATH + + +class DBConfig: + """Database configuration for SQLite.""" + + # Database path - single unified database + DB_PATH: str = os.path.join(RF_DB_PATH, "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 for locked database + DEFAULT_MAX_RETRIES: int = 3 + DEFAULT_BASE_DELAY: float = 0.1 + DEFAULT_MAX_DELAY: float = 1.0 class DatabaseInterface: - """Interface for the database.""" + """Low-level interface for SQLite database operations.""" + + def __init__(self, db_path: str = None): + """ + Initialize database connection. + + Args: + db_path: Path to the database file. Defaults to DBConfig.DB_PATH. + """ + self.db_path = db_path or DBConfig.DB_PATH - def __init__(self): try: - if not os.path.exists(DBConfig.DB_PATH): - path = os.path.dirname(DBConfig.DB_PATH) + # Ensure directory exists + if not os.path.exists(self.db_path): + path = os.path.dirname(self.db_path) os.makedirs(path, exist_ok=True) print(f"Created directory for database at {path}") self.conn: sqlite3.Connection = sqlite3.connect( - DBConfig.DB_PATH, + self.db_path, timeout=DBConfig.CONNECTION_TIMEOUT, check_same_thread=False, isolation_level=None, ) - # Configure database with all PRAGMA settings + # Configure database with performance optimizations pragma_sql = f""" PRAGMA cache_size={DBConfig.CACHE_SIZE}; PRAGMA mmap_size={DBConfig.MMAP_SIZE}; @@ -53,7 +87,7 @@ def retry_on_locked( base_delay: float = DBConfig.DEFAULT_BASE_DELAY, max_delay: float = DBConfig.DEFAULT_MAX_DELAY, ) -> Callable[[Callable[..., Any]], Callable[..., Any]]: - """Decorator to retry operations when database is locked""" + """Decorator to retry operations when database is locked.""" def decorator(func: Callable[..., Any]) -> Callable[..., Any]: @functools.wraps(func) @@ -77,14 +111,16 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: if last_exception: raise last_exception else: - raise RuntimeError("All retries failed but no exception was captured") + raise RuntimeError( + "All retries failed but no exception was captured" + ) return wrapper return decorator def close(self) -> None: - """Close the database connection properly""" + """Close the database connection properly.""" try: if self.conn: self.conn.close() @@ -94,7 +130,7 @@ def close(self) -> None: raise Exception(f"Unexpected error closing database connection: {e}") from e def optimize_periodically(self) -> None: - """Run periodic optimization - call this occasionally, not on every query""" + """Run periodic optimization - call this occasionally, not on every query.""" try: _ = self.conn.execute("PRAGMA optimize") except sqlite3.Error as e: @@ -110,7 +146,22 @@ def execute( fetch: bool = False, commit: bool = False, ) -> list[Any] | tuple[Any] | None: - """Execute a query with automatic retry on database locked errors""" + """ + Execute a query with automatic retry on database locked errors. + + Args: + query: SQL query to execute + params: Query parameters (dict or tuple) + fetch: If True, return fetched results + commit: If True, commit the transaction + + Returns: + Query results if fetch=True, otherwise None + + Raises: + ValueError: If neither fetch nor commit is True + Exception: On database errors + """ # Validate that either fetch or commit is True if not fetch and not commit: raise ValueError("Either fetch or commit must be True") diff --git a/rapidfireai/db/rf_db.py b/rapidfireai/db/rf_db.py new file mode 100644 index 00000000..300a4e43 --- /dev/null +++ b/rapidfireai/db/rf_db.py @@ -0,0 +1,1366 @@ +""" +RapidFire AI Database Manager + +Provides high-level interface for CRUD operations on the experiment database. +Handles experiments, runs, pipelines, contexts, and interactive control operations. +""" + +import json +import os +from typing import Any + +from rapidfireai.db.db_interface import DatabaseInterface +from rapidfireai.utils.constants import ( + ContextStatus, + ControllerTask, + ExperimentStatus, + ExperimentTask, + ICOperation, + ICStatus, + PipelineStatus, + RunEndedBy, + RunSource, + RunStatus, + TaskStatus, + WorkerTask, +) +from rapidfireai.utils.serialize import decode_db_payload, encode_payload + + +class RfDb: + """ + Database manager for RapidFire AI experiments. + + Provides CRUD operations for experiments, runs (training), pipelines (inference), + contexts (RAG), and interactive control operations. + """ + + def __init__(self): + """Initialize the database manager and create tables if needed.""" + self.db = DatabaseInterface() + self._initialize_schema() + + def _initialize_schema(self): + """Initialize database schema from tables.sql file.""" + schema_path = os.path.join(os.path.dirname(__file__), "tables.sql") + if os.path.exists(schema_path): + with open(schema_path) as f: + schema_sql = f.read() + self.db.conn.executescript(schema_sql) + self.db.conn.commit() + + # Run migrations for any schema changes + self._run_migrations() + + def _run_migrations(self): + """Run any necessary schema migrations.""" + # Migration: Add experiments_path to experiments table if it doesn't exist + try: + cursor = self.db.conn.execute("PRAGMA table_info(experiments)") + columns = [row[1] for row in cursor.fetchall()] + + if "experiments_path" not in columns: + self.db.conn.execute("ALTER TABLE experiments ADD COLUMN experiments_path TEXT DEFAULT ''") + self.db.conn.commit() + + if "metric_experiment_id" not in columns: + self.db.conn.execute("ALTER TABLE experiments ADD COLUMN metric_experiment_id TEXT") + self.db.conn.commit() + + if "config" not in columns: + self.db.conn.execute("ALTER TABLE experiments ADD COLUMN config TEXT DEFAULT '{}'") + self.db.conn.commit() + except Exception: + pass + + # Migration: Add metric_run_id to runs table + try: + cursor = self.db.conn.execute("PRAGMA table_info(runs)") + columns = [row[1] for row in cursor.fetchall()] + if "metric_run_id" not in columns: + self.db.conn.execute("ALTER TABLE runs ADD COLUMN metric_run_id TEXT") + self.db.conn.commit() + except Exception: + pass + + # Migration: Add metric_run_id to pipelines table + try: + cursor = self.db.conn.execute("PRAGMA table_info(pipelines)") + columns = [row[1] for row in cursor.fetchall()] + if "metric_run_id" not in columns: + self.db.conn.execute("ALTER TABLE pipelines ADD COLUMN metric_run_id TEXT") + self.db.conn.commit() + except Exception: + pass + + def create_tables(self): + """ + Create database tables (public method for external callers). + Re-runs _initialize_schema() which uses CREATE TABLE IF NOT EXISTS. + """ + self._initialize_schema() + + def close(self): + """Close the database connection.""" + self.db.close() + + # ============================================================================ + # EXPERIMENTS TABLE METHODS (Shared) + # ============================================================================ + + def create_experiment( + self, + experiment_name: str, + experiments_path: str, + metric_experiment_id: str | None = None, + status: ExperimentStatus = ExperimentStatus.RUNNING, + config: dict[str, Any] | None = None, + ) -> int: + """ + Create a new experiment record. + + Args: + experiment_name: Name of the experiment + experiments_path: Path to experiment artifacts directory + metric_experiment_id: Optional MLflow/TensorBoard experiment ID + status: Initial status (default: ExperimentStatus.RUNNING) + config: Optional configuration dict (stored as JSON) + + Returns: + experiment_id of the created experiment + """ + config_json = json.dumps(config) if config else "{}" + + query = """ + INSERT INTO experiments ( + experiment_name, experiments_path, metric_experiment_id, + status, current_task, config, error + ) VALUES (?, ?, ?, ?, ?, ?, '') + """ + self.db.execute( + query, + ( + experiment_name, + experiments_path, + metric_experiment_id, + status.value, + ExperimentTask.IDLE.value, + config_json, + ), + commit=True, + ) + + # Optimize periodically + self.db.optimize_periodically() + + return self.db.cursor.lastrowid + + def get_experiment(self, experiment_id: int) -> dict[str, Any] | None: + """ + Get experiment details by ID. + + Args: + experiment_id: ID of the experiment + + Returns: + Experiment dictionary with all fields, or None if not found + """ + query = """ + SELECT experiment_id, experiment_name, experiments_path, metric_experiment_id, + status, error, created_at, current_task, config + FROM experiments + WHERE experiment_id = ? + """ + result = self.db.execute(query, (experiment_id,), fetch=True) + if result and len(result) > 0: + row = result[0] + return { + "experiment_id": row[0], + "experiment_name": row[1], + "experiments_path": row[2], + "metric_experiment_id": row[3], + "status": ExperimentStatus(row[4]), + "error": row[5], + "created_at": row[6], + "current_task": ExperimentTask(row[7]) if row[7] else ExperimentTask.IDLE, + "config": json.loads(row[8]) if row[8] else {}, + } + return None + + def get_running_experiment(self) -> dict[str, Any]: + """ + Get the currently running experiment (most recent if multiple). + + Returns: + Dictionary with experiment fields + + Raises: + Exception: If no running experiment found + """ + query = """ + SELECT experiment_id, experiment_name, experiments_path, metric_experiment_id, + status, error, created_at, current_task, config + FROM experiments + WHERE status = ? + ORDER BY experiment_id DESC + LIMIT 1 + """ + result = self.db.execute(query, (ExperimentStatus.RUNNING.value,), fetch=True) + if result and len(result) > 0: + row = result[0] + return { + "experiment_id": row[0], + "experiment_name": row[1], + "experiments_path": row[2], + "metric_experiment_id": row[3], + "status": ExperimentStatus(row[4]), + "error": row[5], + "created_at": row[6], + "current_task": ExperimentTask(row[7]) if row[7] else ExperimentTask.IDLE, + "config": json.loads(row[8]) if row[8] else {}, + } + raise Exception("No running experiment found") + + def get_experiment_status(self) -> ExperimentStatus | None: + """Get the status of the most recent experiment.""" + query = """ + SELECT status + FROM experiments + ORDER BY experiment_id DESC + LIMIT 1 + """ + result = self.db.execute(query, fetch=True) + if result: + return ExperimentStatus(result[0][0]) + return None + + 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, experiment_id: int, error: str) -> None: + """Set the error message for an experiment.""" + query = "UPDATE experiments SET error = ? WHERE experiment_id = ?" + self.db.execute(query, (error, experiment_id), commit=True) + + def get_experiment_error(self, experiment_id: int) -> str: + """Get the error message for an experiment.""" + query = "SELECT error FROM experiments WHERE experiment_id = ?" + result = self.db.execute(query, (experiment_id,), fetch=True) + return result[0][0] if result else "" + + def get_all_experiment_names(self) -> list[str]: + """Get all experiment names.""" + query = "SELECT experiment_name FROM experiments" + result = self.db.execute(query, fetch=True) + return [row[0] for row in result] if result else [] + + def get_experiments_path(self, experiment_id_or_name: int | str) -> str: + """Get the experiments path for a given experiment by ID or name.""" + if isinstance(experiment_id_or_name, str): + query = "SELECT experiments_path FROM experiments WHERE experiment_name = ? ORDER BY experiment_id DESC LIMIT 1" + else: + query = "SELECT experiments_path FROM experiments WHERE experiment_id = ?" + result = self.db.execute(query, (experiment_id_or_name,), fetch=True) + if result: + return result[0][0] + raise Exception("Experiments path not found") + + # --- Experiment Config Methods --- + + def get_experiment_config(self, experiment_id: int) -> dict[str, Any]: + """Get the config dict for an experiment.""" + query = "SELECT config FROM experiments WHERE experiment_id = ?" + result = self.db.execute(query, (experiment_id,), fetch=True) + return json.loads(result[0][0]) if result and result[0][0] else {} + + def update_experiment_config(self, experiment_id: int, **kwargs) -> None: + """Update specific keys in the experiment config.""" + config = self.get_experiment_config(experiment_id) + config.update(kwargs) + query = "UPDATE experiments SET config = ? WHERE experiment_id = ?" + self.db.execute(query, (json.dumps(config), experiment_id), commit=True) + + # --- Fit-specific Experiment Methods --- + + def set_experiment_current_task(self, task: ExperimentTask) -> None: + """Set the current task of the running experiment (fit mode).""" + 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 the running experiment (fit mode).""" + query = """ + SELECT current_task + FROM experiments + WHERE status = ? + ORDER BY experiment_id DESC + LIMIT 1 + """ + result = self.db.execute(query, (ExperimentStatus.RUNNING.value,), fetch=True) + if result: + return ExperimentTask(result[0][0]) if result[0][0] else ExperimentTask.IDLE + raise Exception("No running experiment found") + + # --- Evals-specific Experiment Methods --- + + def set_experiment_num_shards(self, experiment_id: int, num_shards: int) -> None: + """Update the number of shards for an experiment (evals mode).""" + self.update_experiment_config(experiment_id, num_shards=num_shards) + + def set_experiment_resources( + self, + experiment_id: int, + num_actors: int, + num_cpus: float = None, + num_gpus: float = None, + ) -> None: + """Update resource allocation for an experiment (evals mode).""" + self.update_experiment_config( + experiment_id, + num_actors=num_actors, + num_cpus=num_cpus, + num_gpus=num_gpus, + ) + + # ============================================================================ + # TABLE RESET METHODS + # ============================================================================ + + def reset_all_tables(self, experiments_table: bool = False) -> None: + """ + Clear data from experiment tables. + + Args: + experiments_table: If True, also clear the experiments table + """ + # Clear tables in order (respecting foreign keys) + tables = [ + "controller_progress", + "worker_progress", + "worker_task", + "actor_tasks", + "interactive_control", + "pipelines", + "contexts", + "runs", + ] + + if experiments_table: + tables.append("experiments") + + for table in tables: + try: + self.db.execute(f"DELETE FROM {table}", commit=True) + except Exception: + pass # Table might not exist + + # Reset auto-increment indices + for table in tables: + try: + self.db.execute("DELETE FROM sqlite_sequence WHERE name = ?", (table,), commit=True) + except Exception: + pass + + def reset_experiment_states(self) -> None: + """ + Reset experiment states when a running task is cancelled. + Marks ongoing tasks as failed for both fit and evals modes. + """ + # Mark all scheduled and in-progress worker tasks as failed (fit) + query = """ + UPDATE worker_task + SET status = ? + WHERE status = ? OR status = ? + """ + try: + self.db.execute( + query, + (TaskStatus.FAILED.value, TaskStatus.IN_PROGRESS.value, TaskStatus.SCHEDULED.value), + commit=True, + ) + except Exception: + pass + + # Mark all scheduled and in-progress actor tasks as failed (evals) + try: + self.db.execute( + query.replace("worker_task", "actor_tasks"), + (TaskStatus.FAILED.value, TaskStatus.IN_PROGRESS.value, TaskStatus.SCHEDULED.value), + commit=True, + ) + except Exception: + pass + + # Mark ongoing and new runs as failed (fit) + try: + run_query = """ + UPDATE runs + SET status = ? + WHERE status = ? OR status = ? + """ + self.db.execute( + run_query, + (RunStatus.FAILED.value, RunStatus.ONGOING.value, RunStatus.NEW.value), + commit=True, + ) + except Exception: + pass + + # Mark ongoing and new pipelines as failed (evals) + try: + pipeline_query = """ + UPDATE pipelines + SET status = ? + WHERE status = ? OR status = ? + """ + self.db.execute( + pipeline_query, + (PipelineStatus.FAILED.value, PipelineStatus.ONGOING.value, PipelineStatus.NEW.value), + commit=True, + ) + except Exception: + pass + + # Mark ongoing and new contexts as failed (evals) + try: + context_query = """ + UPDATE contexts + SET status = ? + WHERE status = ? OR status = ? + """ + self.db.execute( + context_query, + (ContextStatus.FAILED.value, ContextStatus.ONGOING.value, ContextStatus.NEW.value), + commit=True, + ) + except Exception: + pass + + # Reset all pending interactive control tasks + try: + ic_query = """ + UPDATE interactive_control + SET status = ? + WHERE status = ? + """ + self.db.execute( + ic_query, + (ICStatus.SKIPPED.value, ICStatus.PENDING.value), + commit=True, + ) + except Exception: + pass + + # Reset progress tables (fit) + for table in ["controller_progress", "worker_progress"]: + try: + self.db.execute(f"DELETE FROM {table}", commit=True) + except Exception: + pass + + # ============================================================================ + # INTERACTIVE CONTROL TABLE METHODS (Unified) + # ============================================================================ + + def create_ic_operation( + self, + target_id: int, + target_type: str, + operation: str | ICOperation, + config_data: dict[str, Any] | str | None = None, + ) -> int: + """ + Create a new interactive control operation. + + Args: + target_id: run_id (fit) or pipeline_id (evals) + target_type: 'run' or 'pipeline' + operation: Operation type (stop, resume, delete, clone, clone_warm) + config_data: JSON config (config_leaf for fit, request_data for evals) + + Returns: + ic_id of the created operation + """ + if isinstance(operation, ICOperation): + operation = operation.value + + if isinstance(config_data, dict): + config_data = json.dumps(config_data) + elif config_data is None: + config_data = "{}" + + query = """ + INSERT INTO interactive_control ( + target_id, target_type, operation, config_data, status + ) VALUES (?, ?, ?, ?, ?) + """ + self.db.execute( + query, + (target_id, target_type, operation, config_data, ICStatus.PENDING.value), + commit=True, + ) + return self.db.cursor.lastrowid + + def get_pending_ic_operations(self, target_type: str = None) -> list[dict[str, Any]]: + """ + Get all pending IC operations. + + Args: + target_type: Optional filter for 'run' or 'pipeline' + + Returns: + List of IC operation dictionaries + """ + if target_type: + query = """ + SELECT ic_id, target_id, target_type, operation, config_data, status, error, created_at, processed_at + FROM interactive_control + WHERE status = ? AND target_type = ? + ORDER BY created_at ASC + """ + results = self.db.execute(query, (ICStatus.PENDING.value, target_type), fetch=True) + else: + query = """ + SELECT ic_id, target_id, target_type, operation, config_data, status, error, created_at, processed_at + FROM interactive_control + WHERE status = ? + ORDER BY created_at ASC + """ + results = self.db.execute(query, (ICStatus.PENDING.value,), fetch=True) + + operations = [] + for row in results or []: + operations.append({ + "ic_id": row[0], + "target_id": row[1], + "target_type": row[2], + "operation": row[3], + "config_data": json.loads(row[4]) if row[4] else {}, + "status": row[5], + "error": row[6], + "created_at": row[7], + "processed_at": row[8], + }) + return operations + + def get_ic_operation(self, ic_id: int) -> dict[str, Any] | None: + """Get IC operation by ID.""" + query = """ + SELECT ic_id, target_id, target_type, operation, config_data, status, error, created_at, processed_at + FROM interactive_control + WHERE ic_id = ? + """ + result = self.db.execute(query, (ic_id,), fetch=True) + if result: + row = result[0] + return { + "ic_id": row[0], + "target_id": row[1], + "target_type": row[2], + "operation": row[3], + "config_data": json.loads(row[4]) if row[4] else {}, + "status": row[5], + "error": row[6], + "created_at": row[7], + "processed_at": row[8], + } + return None + + def update_ic_operation_status(self, ic_id: int, status: str | ICStatus, error: str = "") -> None: + """Update IC operation status.""" + import time + + if isinstance(status, ICStatus): + status = status.value + + query = """ + UPDATE interactive_control + SET status = ?, error = ?, processed_at = ? + WHERE ic_id = ? + """ + self.db.execute(query, (status, error, time.time(), ic_id), commit=True) + + def get_all_ic_operations(self) -> list[dict[str, Any]]: + """Get all IC operations.""" + query = """ + SELECT ic_id, target_id, target_type, operation, config_data, status, error, created_at, processed_at + FROM interactive_control + ORDER BY created_at DESC + """ + results = self.db.execute(query, fetch=True) + operations = [] + for row in results or []: + operations.append({ + "ic_id": row[0], + "target_id": row[1], + "target_type": row[2], + "operation": row[3], + "config_data": json.loads(row[4]) if row[4] else {}, + "status": row[5], + "error": row[6], + "created_at": row[7], + "processed_at": row[8], + }) + return operations + + # ============================================================================ + # FIT MODE: RUNS TABLE METHODS + # ============================================================================ + + def create_run( + self, + config_leaf: dict[str, Any], + status: RunStatus, + metric_run_id: str | None = None, + flattened_config: dict[str, Any] | None = None, + completed_steps: int = 0, + total_steps: int = 0, + num_chunks_visited_curr_epoch: int = 0, + num_epochs_completed: int = 0, + chunk_offset: int = 0, + error: str = "", + source: RunSource | None = None, + ended_by: RunEndedBy | None = None, + warm_started_from: int | None = None, + cloned_from: int | None = None, + estimated_runtime: float = 0.0, + required_workers: int = 0, + ) -> int: + """Create a new run (fit mode).""" + query = """ + INSERT INTO runs ( + status, metric_run_id, flattened_config, config_leaf, + completed_steps, total_steps, num_chunks_visited_curr_epoch, + num_epochs_completed, chunk_offset, error, source, ended_by, + warm_started_from, cloned_from, estimated_runtime, required_workers + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """ + self.db.execute( + query, + ( + status.value, + metric_run_id, + json.dumps(flattened_config) if flattened_config else "{}", + encode_payload(config_leaf) if config_leaf else "{}", + completed_steps, + total_steps, + num_chunks_visited_curr_epoch, + num_epochs_completed, + chunk_offset, + error, + source.value if source else "", + ended_by.value if ended_by else "", + warm_started_from, + cloned_from, + estimated_runtime, + required_workers, + ), + commit=True, + ) + result = self.db.execute("SELECT last_insert_rowid()", fetch=True) + if result: + return result[0][0] + raise Exception("Failed to create run") + + def get_run(self, run_id: int) -> dict[str, Any]: + """Get a run's details (fit mode).""" + query = """ + SELECT status, metric_run_id, flattened_config, config_leaf, + completed_steps, total_steps, num_chunks_visited_curr_epoch, + num_epochs_completed, chunk_offset, error, source, ended_by, + warm_started_from, cloned_from, estimated_runtime, required_workers + FROM runs + WHERE run_id = ? + """ + result = self.db.execute(query, (run_id,), fetch=True) + if result: + row = result[0] + return { + "status": RunStatus(row[0]), + "metric_run_id": row[1], + "flattened_config": json.loads(row[2]) if row[2] else {}, + "config_leaf": decode_db_payload(row[3]) if row[3] and row[3] != "{}" else {}, + "completed_steps": row[4], + "total_steps": row[5], + "num_chunks_visited_curr_epoch": row[6], + "num_epochs_completed": row[7], + "chunk_offset": row[8], + "error": row[9], + "source": RunSource(row[10]) if row[10] else None, + "ended_by": RunEndedBy(row[11]) if row[11] else None, + "warm_started_from": row[12], + "cloned_from": row[13], + "estimated_runtime": row[14], + "required_workers": row[15], + } + raise Exception("Run not found") + + def get_all_runs(self) -> dict[int, dict[str, Any]]: + """Get all runs (fit mode).""" + query = """ + SELECT run_id, status, metric_run_id, flattened_config, config_leaf, + completed_steps, total_steps, num_chunks_visited_curr_epoch, + num_epochs_completed, chunk_offset, error, source, ended_by, + warm_started_from, cloned_from, estimated_runtime, required_workers + FROM runs + """ + results = self.db.execute(query, fetch=True) + formatted = {} + if results: + for row in results: + formatted[row[0]] = { + "status": RunStatus(row[1]), + "metric_run_id": row[2], + "flattened_config": json.loads(row[3]) if row[3] else {}, + "config_leaf": decode_db_payload(row[4]) if row[4] and row[4] != "{}" else {}, + "completed_steps": row[5], + "total_steps": row[6], + "num_chunks_visited_curr_epoch": row[7], + "num_epochs_completed": row[8], + "chunk_offset": row[9], + "error": row[10], + "source": RunSource(row[11]) if row[11] else None, + "ended_by": RunEndedBy(row[12]) if row[12] else None, + "warm_started_from": row[13], + "cloned_from": row[14], + "estimated_runtime": row[15], + "required_workers": row[16], + } + return formatted + + def get_runs_by_status(self, statuses: list[RunStatus]) -> dict[int, dict[str, Any]]: + """Get all runs by statuses (fit mode).""" + if not statuses: + return {} + + placeholders = ",".join(["?"] * len(statuses)) + query = f""" + SELECT run_id, status, metric_run_id, flattened_config, config_leaf, + completed_steps, total_steps, num_chunks_visited_curr_epoch, + num_epochs_completed, chunk_offset, error, source, ended_by, + warm_started_from, cloned_from, estimated_runtime, required_workers + FROM runs + WHERE status IN ({placeholders}) + """ + status_values = [s.value for s in statuses] + results = self.db.execute(query, status_values, fetch=True) + formatted = {} + if results: + for row in results: + formatted[row[0]] = { + "status": RunStatus(row[1]), + "metric_run_id": row[2], + "flattened_config": json.loads(row[3]) if row[3] else {}, + "config_leaf": decode_db_payload(row[4]) if row[4] and row[4] != "{}" else {}, + "completed_steps": row[5], + "total_steps": row[6], + "num_chunks_visited_curr_epoch": row[7], + "num_epochs_completed": row[8], + "chunk_offset": row[9], + "error": row[10], + "source": RunSource(row[11]) if row[11] else None, + "ended_by": RunEndedBy(row[12]) if row[12] else None, + "warm_started_from": row[13], + "cloned_from": row[14], + "estimated_runtime": row[15], + "required_workers": row[16], + } + return formatted + + def set_run_status(self, run_id: int, status: RunStatus) -> None: + """Set the status of a run (fit mode).""" + query = "UPDATE runs SET status = ? WHERE run_id = ?" + self.db.execute(query, (status.value, run_id), commit=True) + + def set_run_details( + self, + run_id: int, + status: RunStatus | None = None, + metric_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, + num_chunks_visited_curr_epoch: int | None = None, + num_epochs_completed: int | None = None, + chunk_offset: int | None = None, + error: str | None = None, + source: RunSource | str | None = None, + ended_by: RunEndedBy | str | None = None, + warm_started_from: int | None = None, + cloned_from: int | None = None, + estimated_runtime: float | None = None, + required_workers: int | None = None, + ) -> None: + """Set details of an existing run (fit mode).""" + columns = { + "status": status.value if isinstance(status, RunStatus) else status, + "metric_run_id": metric_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, + "num_chunks_visited_curr_epoch": num_chunks_visited_curr_epoch, + "num_epochs_completed": num_epochs_completed, + "chunk_offset": chunk_offset, + "error": error, + "source": source.value if isinstance(source, RunSource) else source, + "ended_by": ended_by.value if isinstance(ended_by, RunEndedBy) else ended_by, + "warm_started_from": warm_started_from, + "cloned_from": cloned_from, + "estimated_runtime": estimated_runtime, + "required_workers": required_workers, + } + + columns = {k: v for k, v in columns.items() if v is not None} + if not columns: + return + + query_parts = [f"{col} = ?" for col in columns] + values = list(columns.values()) + values.append(run_id) + + query = f"UPDATE runs SET {', '.join(query_parts)} WHERE run_id = ?" + self.db.execute(query, tuple(values), commit=True) + + def set_completed_steps(self, run_id: int, completed_steps: int) -> None: + """Set the completed steps for a run (fit mode).""" + query = "UPDATE runs SET completed_steps = ? WHERE run_id = ?" + self.db.execute(query, (completed_steps, run_id), commit=True) + + def set_estimated_runtime(self, run_id: int, estimated_runtime: float) -> None: + """Set the estimated runtime per batch for a run (fit mode).""" + query = "UPDATE runs SET estimated_runtime = ? WHERE run_id = ?" + self.db.execute(query, (estimated_runtime, run_id), commit=True) + + def get_completed_steps(self, run_id: int) -> int: + """Get the completed steps for a run (fit mode).""" + query = "SELECT completed_steps FROM runs WHERE run_id = ?" + result = self.db.execute(query, (run_id,), fetch=True) + if result: + return result[0][0] + raise Exception("Run not found") + + # ============================================================================ + # FIT MODE: WORKER TASK TABLE METHODS + # ============================================================================ + + def create_worker_task( + self, + worker_id: int, + task_type: WorkerTask, + status: TaskStatus, + run_id: int, + chunk_id: int = -1, + multi_worker_details: dict[str, Any] | None = None, + config_options: dict[str, Any] | None = None, + ) -> int: + """Create a worker task (fit mode).""" + query = """ + INSERT INTO worker_task (worker_id, task_type, status, run_id, chunk_id, multi_worker_details, config_options) + VALUES (?, ?, ?, ?, ?, ?, ?) + """ + multi_worker_details_str = json.dumps(multi_worker_details) if multi_worker_details else "{}" + 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, multi_worker_details_str, config_options_str), + commit=True, + ) + result = self.db.execute("SELECT last_insert_rowid()", fetch=True) + if result: + return result[0][0] + raise Exception("Failed to create worker task") + + def get_worker_scheduled_task(self, worker_id: int) -> dict[str, Any]: + """Get the latest scheduled task for a worker (fit mode).""" + query = """ + SELECT task_id, task_type, run_id, chunk_id, multi_worker_details, config_options + FROM worker_task + WHERE worker_id = ? AND status = ? + ORDER BY task_id DESC + LIMIT 1 + """ + result = self.db.execute(query, (worker_id, TaskStatus.SCHEDULED.value), fetch=True) + if result: + row = result[0] + multi_worker_details = json.loads(row[4]) if row[4] and row[4] != "{}" else {} + config_options = decode_db_payload(row[5]) if row[5] and row[5] != "{}" else {} + return { + "task_id": row[0], + "task_type": WorkerTask(row[1]), + "run_id": row[2], + "chunk_id": row[3], + "multi_worker_details": multi_worker_details, + "config_options": config_options, + } + return {} + + def get_all_worker_tasks(self) -> dict[int, dict[str, Any]]: + """Get the latest task of each worker (fit mode).""" + query = """ + SELECT worker_id, task_id, task_type, status, run_id, chunk_id, multi_worker_details, config_options + FROM worker_task wt1 + WHERE task_id = ( + SELECT MAX(task_id) + FROM worker_task wt2 + WHERE wt2.worker_id = wt1.worker_id + ) + """ + results = self.db.execute(query, fetch=True) + formatted = {} + if results: + for row in results: + formatted[row[0]] = { + "task_id": row[1], + "task_type": WorkerTask(row[2]), + "status": TaskStatus(row[3]), + "run_id": row[4], + "chunk_id": row[5], + "multi_worker_details": json.loads(row[6]) if row[6] and row[6] != "{}" else {}, + "config_options": decode_db_payload(row[7]) if row[7] and row[7] != "{}" else {}, + } + return formatted + + def set_worker_task_status(self, worker_id: int, status: TaskStatus) -> None: + """Set the status of the latest task of a worker (fit mode).""" + 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) + + def cancel_worker_tasks_for_run(self, run_id: int) -> None: + """Cancel all scheduled worker tasks for a given run.""" + query = """ + UPDATE worker_task + SET status = ? + WHERE run_id = ? AND status = ? + """ + self.db.execute(query, (TaskStatus.SKIPPED.value, run_id, TaskStatus.SCHEDULED.value), commit=True) + + # ============================================================================ + # FIT MODE: PROGRESS TABLE METHODS + # ============================================================================ + + def set_controller_progress(self, run_id: int, progress: float) -> None: + """Set the controller progress for a run (fit mode).""" + query = """ + INSERT INTO controller_progress (run_id, progress) + VALUES (?, ?) + ON CONFLICT (run_id) + DO UPDATE SET progress = EXCLUDED.progress + """ + self.db.execute(query, (run_id, round(progress, 2)), commit=True) + + def get_controller_progress(self, run_id: int) -> float: + """Get the controller progress for a run (fit mode).""" + query = "SELECT progress FROM controller_progress WHERE run_id = ?" + result = self.db.execute(query, (run_id,), fetch=True) + return result[0][0] if result else 0.0 + + def set_worker_progress(self, run_id: int, subchunk_progress: float) -> None: + """Set the worker progress for a run (fit mode).""" + query = """ + INSERT INTO worker_progress (run_id, subchunk_progress) + VALUES (?, ?) + ON CONFLICT (run_id) + DO UPDATE SET subchunk_progress = EXCLUDED.subchunk_progress + """ + self.db.execute(query, (run_id, round(subchunk_progress, 2)), commit=True) + + def get_worker_progress(self, run_id: int) -> float: + """Get the worker progress for a run (fit mode).""" + query = "SELECT subchunk_progress FROM worker_progress WHERE run_id = ?" + result = self.db.execute(query, (run_id,), fetch=True) + return result[0][0] if result else 0.0 + + # ============================================================================ + # EVALS MODE: CONTEXTS TABLE METHODS + # ============================================================================ + + def create_context( + self, + context_hash: str, + rag_config_json: str = None, + prompt_config_json: str = None, + status: ContextStatus = ContextStatus.NEW, + ) -> int: + """Create a new context record (evals mode).""" + # Check if context with this hash already exists + query = "SELECT context_id FROM contexts WHERE context_hash = ?" + result = self.db.execute(query, (context_hash,), fetch=True) + if result: + return result[0][0] + + query = """ + INSERT INTO contexts (context_hash, rag_config_json, prompt_config_json, status, error) + VALUES (?, ?, ?, ?, '') + """ + self.db.execute(query, (context_hash, rag_config_json, prompt_config_json, status.value), commit=True) + return self.db.cursor.lastrowid + + def get_context(self, context_id: int) -> dict[str, Any] | None: + """Get context by ID (evals mode).""" + query = """ + SELECT context_id, context_hash, rag_config_json, prompt_config_json, + status, error, started_at, completed_at, duration_seconds + FROM contexts + WHERE context_id = ? + """ + result = self.db.execute(query, (context_id,), fetch=True) + if result: + row = result[0] + return { + "context_id": row[0], + "context_hash": row[1], + "rag_config_json": row[2], + "prompt_config_json": row[3], + "status": ContextStatus(row[4]), + "error": row[5], + "started_at": row[6], + "completed_at": row[7], + "duration_seconds": row[8], + } + return None + + def get_context_by_hash(self, context_hash: str) -> dict[str, Any] | None: + """Get context by hash (evals mode).""" + query = """ + SELECT context_id, context_hash, rag_config_json, prompt_config_json, + status, error, started_at, completed_at, duration_seconds + FROM contexts + WHERE context_hash = ? + """ + result = self.db.execute(query, (context_hash,), fetch=True) + if result: + row = result[0] + return { + "context_id": row[0], + "context_hash": row[1], + "rag_config_json": row[2], + "prompt_config_json": row[3], + "status": ContextStatus(row[4]), + "error": row[5], + "started_at": row[6], + "completed_at": row[7], + "duration_seconds": row[8], + } + return None + + def set_context_status(self, context_id: int, status: ContextStatus) -> None: + """Update context status (evals mode).""" + query = "UPDATE contexts SET status = ? WHERE context_id = ?" + self.db.execute(query, (status.value, context_id), commit=True) + + def set_context_start_time(self, context_id: int, start_time: str) -> None: + """Set start time for context building (evals mode).""" + query = "UPDATE contexts SET started_at = ? WHERE context_id = ?" + self.db.execute(query, (start_time, context_id), commit=True) + + def set_context_end_time(self, context_id: int, end_time: str, duration_seconds: float) -> None: + """Set end time and duration for context building (evals mode).""" + query = "UPDATE contexts SET completed_at = ?, duration_seconds = ? WHERE context_id = ?" + self.db.execute(query, (end_time, duration_seconds, context_id), commit=True) + + def set_context_error(self, context_id: int, error: str) -> None: + """Set error message for a context (evals mode).""" + query = "UPDATE contexts SET error = ? WHERE context_id = ?" + self.db.execute(query, (error, context_id), commit=True) + + # ============================================================================ + # EVALS MODE: PIPELINES TABLE METHODS + # ============================================================================ + + def create_pipeline( + self, + pipeline_type: str, + pipeline_config: Any, + context_id: int = None, + status: PipelineStatus = PipelineStatus.NEW, + flattened_config: dict[str, Any] = None, + ) -> int: + """Create a new pipeline record (evals mode).""" + encoded_config = encode_payload(pipeline_config) + + # Extract JSON-serializable data + from rapidfireai.utils.serialize import extract_pipeline_config_json + json_config_dict = extract_pipeline_config_json(pipeline_config) + json_config_str = json.dumps(json_config_dict) if json_config_dict else "{}" + flattened_config_str = json.dumps(flattened_config) if flattened_config else "{}" + + query = """ + INSERT INTO pipelines ( + context_id, pipeline_type, pipeline_config, pipeline_config_json, + flattened_config, status, error, current_shard_id, shards_completed, + total_samples_processed, metric_run_id + ) VALUES (?, ?, ?, ?, ?, ?, '', 0, 0, 0, NULL) + """ + self.db.execute( + query, + (context_id, pipeline_type, encoded_config, json_config_str, flattened_config_str, status.value), + commit=True, + ) + return self.db.cursor.lastrowid + + def get_pipeline(self, pipeline_id: int | str) -> dict[str, Any] | None: + """Get a single pipeline by ID (evals mode).""" + # Try as metric_run_id first if it's a string + if isinstance(pipeline_id, str): + pipeline = self.get_pipeline_by_metric_run_id(pipeline_id) + if pipeline: + return pipeline + + query = """ + SELECT pipeline_id, context_id, pipeline_type, pipeline_config, + pipeline_config_json, flattened_config, status, current_shard_id, + shards_completed, total_samples_processed, metric_run_id, error, created_at + FROM pipelines + WHERE pipeline_id = ? + """ + result = self.db.execute(query, (pipeline_id,), fetch=True) + if result: + row = result[0] + return { + "pipeline_id": row[0], + "context_id": row[1], + "pipeline_type": row[2], + "pipeline_config": decode_db_payload(row[3]) if row[3] else None, + "pipeline_config_json": json.loads(row[4]) if row[4] else None, + "flattened_config": json.loads(row[5]) if row[5] else {}, + "status": PipelineStatus(row[6]), + "current_shard_id": row[7], + "shards_completed": row[8], + "total_samples_processed": row[9], + "metric_run_id": row[10], + "error": row[11], + "created_at": row[12], + } + return None + + def get_pipeline_by_metric_run_id(self, metric_run_id: str) -> dict[str, Any] | None: + """Get pipeline by its metric_run_id (evals mode).""" + query = """ + SELECT pipeline_id, context_id, pipeline_type, pipeline_config, + pipeline_config_json, flattened_config, status, current_shard_id, + shards_completed, total_samples_processed, metric_run_id, error, created_at + FROM pipelines + WHERE metric_run_id = ? + """ + result = self.db.execute(query, (metric_run_id,), fetch=True) + if result: + row = result[0] + return { + "pipeline_id": row[0], + "context_id": row[1], + "pipeline_type": row[2], + "pipeline_config": decode_db_payload(row[3]) if row[3] else None, + "pipeline_config_json": json.loads(row[4]) if row[4] else None, + "flattened_config": json.loads(row[5]) if row[5] else {}, + "status": PipelineStatus(row[6]), + "current_shard_id": row[7], + "shards_completed": row[8], + "total_samples_processed": row[9], + "metric_run_id": row[10], + "error": row[11], + "created_at": row[12], + } + return None + + def get_all_pipelines(self) -> list[dict[str, Any]]: + """Get all pipelines (evals mode).""" + query = """ + SELECT pipeline_id, context_id, pipeline_type, pipeline_config, + pipeline_config_json, flattened_config, status, current_shard_id, + shards_completed, total_samples_processed, metric_run_id, error, created_at + FROM pipelines + ORDER BY pipeline_id DESC + """ + results = self.db.execute(query, fetch=True) + pipelines = [] + if results: + for row in results: + pipelines.append({ + "pipeline_id": row[0], + "context_id": row[1], + "pipeline_type": row[2], + "pipeline_config": decode_db_payload(row[3]) if row[3] else None, + "pipeline_config_json": json.loads(row[4]) if row[4] else None, + "flattened_config": json.loads(row[5]) if row[5] else {}, + "status": PipelineStatus(row[6]), + "current_shard_id": row[7], + "shards_completed": row[8], + "total_samples_processed": row[9], + "metric_run_id": row[10], + "error": row[11], + "created_at": row[12], + }) + return pipelines + + def get_all_pipeline_ids(self) -> list[dict[str, Any]]: + """Get lightweight list of pipeline IDs with minimal info (evals mode).""" + query = """ + SELECT pipeline_id, status, shards_completed, total_samples_processed + FROM pipelines + ORDER BY pipeline_id DESC + """ + results = self.db.execute(query, fetch=True) + pipelines = [] + if results: + for row in results: + pipelines.append({ + "pipeline_id": row[0], + "status": row[1], + "shards_completed": row[2], + "total_samples_processed": row[3], + }) + return pipelines + + def get_pipeline_config_json(self, pipeline_id: int) -> dict[str, Any] | None: + """Get only the config JSON for a specific pipeline (evals mode).""" + query = "SELECT pipeline_config_json, context_id FROM pipelines WHERE pipeline_id = ?" + result = self.db.execute(query, (pipeline_id,), fetch=True) + if result and result[0][0]: + return { + "pipeline_config_json": json.loads(result[0][0]), + "context_id": result[0][1], + } + return None + + def set_pipeline_status(self, pipeline_id: int, status: PipelineStatus) -> None: + """Update pipeline status (evals mode).""" + query = "UPDATE pipelines SET status = ? WHERE pipeline_id = ?" + self.db.execute(query, (status.value, pipeline_id), commit=True) + + def set_pipeline_progress( + self, + pipeline_id: int, + current_shard_id: int, + shards_completed: int, + total_samples_processed: int, + ) -> None: + """Update pipeline progress metrics (evals mode).""" + query = """ + UPDATE pipelines + SET current_shard_id = ?, shards_completed = ?, total_samples_processed = ? + WHERE pipeline_id = ? + """ + self.db.execute(query, (current_shard_id, shards_completed, total_samples_processed, pipeline_id), commit=True) + + def set_pipeline_current_shard(self, pipeline_id: int, shard_id: int) -> None: + """Update the current shard being processed by a pipeline (evals mode).""" + query = "UPDATE pipelines SET current_shard_id = ? WHERE pipeline_id = ?" + self.db.execute(query, (shard_id, pipeline_id), commit=True) + + def set_pipeline_error(self, pipeline_id: int, error: str) -> None: + """Set error message for a pipeline (evals mode).""" + query = "UPDATE pipelines SET error = ? WHERE pipeline_id = ?" + self.db.execute(query, (error, pipeline_id), commit=True) + + def set_pipeline_metric_run_id(self, pipeline_id: int, metric_run_id: str) -> None: + """Set MetricLogger run ID for a pipeline (evals mode).""" + query = "UPDATE pipelines SET metric_run_id = ? WHERE pipeline_id = ?" + self.db.execute(query, (metric_run_id, pipeline_id), commit=True) + + # ============================================================================ + # EVALS MODE: ACTOR TASKS TABLE METHODS + # ============================================================================ + + def create_actor_task( + self, + pipeline_id: int, + actor_id: int, + shard_id: int, + status: TaskStatus = TaskStatus.SCHEDULED, + ) -> int: + """Create a new actor task record (evals mode).""" + query = """ + INSERT INTO actor_tasks (pipeline_id, actor_id, shard_id, status, error_message) + VALUES (?, ?, ?, ?, '') + """ + self.db.execute(query, (pipeline_id, actor_id, shard_id, status.value), commit=True) + return self.db.cursor.lastrowid + + def get_actor_task(self, task_id: int) -> dict[str, Any] | None: + """Get actor task by ID (evals mode).""" + query = """ + SELECT task_id, pipeline_id, actor_id, shard_id, status, + error_message, started_at, completed_at, duration_seconds + FROM actor_tasks + WHERE task_id = ? + """ + result = self.db.execute(query, (task_id,), fetch=True) + if result: + row = result[0] + return { + "task_id": row[0], + "pipeline_id": row[1], + "actor_id": row[2], + "shard_id": row[3], + "status": TaskStatus(row[4]), + "error_message": row[5], + "started_at": row[6], + "completed_at": row[7], + "duration_seconds": row[8], + } + return None + + def get_running_actor_tasks(self) -> list[dict[str, Any]]: + """Get all currently running actor tasks (evals mode).""" + query = """ + SELECT task_id, pipeline_id, actor_id, shard_id, status, + error_message, started_at, completed_at, duration_seconds + FROM actor_tasks + WHERE status = ? + ORDER BY task_id DESC + """ + results = self.db.execute(query, (TaskStatus.IN_PROGRESS.value,), fetch=True) + tasks = [] + if results: + for row in results: + tasks.append({ + "task_id": row[0], + "pipeline_id": row[1], + "actor_id": row[2], + "shard_id": row[3], + "status": TaskStatus(row[4]), + "error_message": row[5], + "started_at": row[6], + "completed_at": row[7], + "duration_seconds": row[8], + }) + return tasks + + def set_actor_task_status(self, task_id: int, status: TaskStatus) -> None: + """Update actor task status (evals mode).""" + query = "UPDATE actor_tasks SET status = ? WHERE task_id = ?" + self.db.execute(query, (status.value, task_id), commit=True) + + def set_actor_task_start_time(self, task_id: int, start_time: str) -> None: + """Set start time for an actor task (evals mode).""" + query = "UPDATE actor_tasks SET started_at = ? WHERE task_id = ?" + self.db.execute(query, (start_time, task_id), commit=True) + + def set_actor_task_end_time(self, task_id: int, end_time: str, duration_seconds: float) -> None: + """Set end time and duration for an actor task (evals mode).""" + query = "UPDATE actor_tasks SET completed_at = ?, duration_seconds = ? WHERE task_id = ?" + self.db.execute(query, (end_time, duration_seconds, task_id), commit=True) + + def set_actor_task_error(self, task_id: int, error_message: str) -> None: + """Set error message for an actor task (evals mode).""" + query = "UPDATE actor_tasks SET error_message = ? WHERE task_id = ?" + self.db.execute(query, (error_message, task_id), commit=True) + + +# Backwards compatibility aliases +RFDatabase = RfDb + +__all__ = ["RfDb", "RFDatabase", "encode_payload", "decode_db_payload"] diff --git a/rapidfireai/db/tables.sql b/rapidfireai/db/tables.sql new file mode 100644 index 00000000..d9d3f51e --- /dev/null +++ b/rapidfireai/db/tables.sql @@ -0,0 +1,144 @@ +-- RapidFire AI Unified Database Schema +-- SQLite schema for tracking both fit (training) and evals (inference) experiments +-- Unified: Single experiments table, mode-specific tables for runs/pipelines + +-- ============================================================================ +-- EXPERIMENTS TABLE (Shared) +-- Tracks high-level experiment configuration and status for both modes +-- ============================================================================ +CREATE TABLE IF NOT EXISTS experiments ( + experiment_id INTEGER PRIMARY KEY AUTOINCREMENT, + experiment_name TEXT NOT NULL, + experiments_path TEXT NOT NULL, + metric_experiment_id TEXT, + status TEXT NOT NULL, -- 'running', 'completed', 'failed', 'cancelled' + error TEXT DEFAULT '', + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + + -- Fit-specific state (nullable, not used by evals) + current_task TEXT DEFAULT 'idle', -- 'idle', 'create_models', 'run_fit', 'ic_ops' + + -- Mode-specific configuration (JSON) + config TEXT DEFAULT '{}' +); + +-- ============================================================================ +-- INTERACTIVE_CONTROL TABLE (Unified) +-- Tracks user-initiated dynamic control operations for both runs and pipelines +-- ============================================================================ +CREATE TABLE IF NOT EXISTS interactive_control ( + ic_id INTEGER PRIMARY KEY AUTOINCREMENT, + target_id INTEGER NOT NULL, -- run_id for fit, pipeline_id for evals + target_type TEXT NOT NULL, -- 'run' or 'pipeline' + operation TEXT NOT NULL, -- 'stop', 'resume', 'delete', 'clone', 'clone_warm' + config_data TEXT DEFAULT '{}', -- JSON: config_leaf for fit, request_data for evals + status TEXT NOT NULL, -- 'pending', 'processing', 'completed', 'failed', 'skipped' + error TEXT DEFAULT '', + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + processed_at TIMESTAMP +); + +-- ============================================================================ +-- FIT MODE TABLES +-- Tables for training experiment tracking +-- ============================================================================ + +-- Runs table (fit mode) +CREATE TABLE IF NOT EXISTS runs ( + run_id INTEGER PRIMARY KEY AUTOINCREMENT, + status TEXT NOT NULL, + metric_run_id TEXT, + flattened_config TEXT DEFAULT '{}', + config_leaf TEXT DEFAULT '{}', + completed_steps INTEGER DEFAULT 0, + total_steps INTEGER DEFAULT 0, + num_chunks_visited_curr_epoch INTEGER DEFAULT 0, + num_epochs_completed INTEGER DEFAULT 0, + chunk_offset INTEGER DEFAULT 0, + error TEXT DEFAULT '', + source TEXT DEFAULT '', + ended_by TEXT DEFAULT '', + warm_started_from INTEGER DEFAULT NULL, + cloned_from INTEGER DEFAULT NULL, + estimated_runtime REAL DEFAULT 0.0, + required_workers INTEGER DEFAULT 0 +); + +-- Worker Task table (fit mode) +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, + multi_worker_details TEXT DEFAULT '{}', + config_options TEXT DEFAULT '{}', + FOREIGN KEY (run_id) REFERENCES runs (run_id) +); + +-- Controller Progress table (fit mode) +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 (fit mode) +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) +); + +-- ============================================================================ +-- EVALS MODE TABLES +-- Tables for inference experiment tracking +-- ============================================================================ + +-- Contexts table (evals mode - RAG context configurations) +CREATE TABLE IF NOT EXISTS contexts ( + context_id INTEGER PRIMARY KEY AUTOINCREMENT, + context_hash TEXT NOT NULL UNIQUE, -- SHA256 hash for deduplication + rag_config_json TEXT, + prompt_config_json TEXT, + status TEXT NOT NULL, -- 'new', 'ongoing', 'completed', 'deleted', 'failed' + error TEXT DEFAULT '', + started_at TIMESTAMP, + completed_at TIMESTAMP, + duration_seconds REAL +); + +-- Pipelines table (evals mode) +CREATE TABLE IF NOT EXISTS pipelines ( + pipeline_id INTEGER PRIMARY KEY AUTOINCREMENT, + context_id INTEGER, + pipeline_type TEXT NOT NULL, -- 'vllm', 'openai_api', etc. + pipeline_config TEXT NOT NULL, -- Encoded with dill (includes functions/classes) + pipeline_config_json TEXT, -- JSON-serializable version for display + flattened_config TEXT DEFAULT '{}', + status TEXT NOT NULL, -- 'new', 'ongoing', 'completed', 'stopped', 'deleted', 'failed' + current_shard_id INTEGER DEFAULT 0, + shards_completed INTEGER DEFAULT 0, + total_samples_processed INTEGER DEFAULT 0, + metric_run_id TEXT, + error TEXT DEFAULT '', + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + + FOREIGN KEY (context_id) REFERENCES contexts(context_id) ON DELETE SET NULL +); + +-- Actor Tasks table (evals mode) +CREATE TABLE IF NOT EXISTS actor_tasks ( + task_id INTEGER PRIMARY KEY AUTOINCREMENT, + pipeline_id INTEGER NOT NULL, + actor_id INTEGER NOT NULL, + shard_id INTEGER NOT NULL, + status TEXT NOT NULL, -- 'scheduled', 'in_progress', 'completed', 'failed' + error_message TEXT DEFAULT '', + started_at TIMESTAMP, + completed_at TIMESTAMP, + duration_seconds REAL, + + FOREIGN KEY (pipeline_id) REFERENCES pipelines(pipeline_id) ON DELETE CASCADE +); diff --git a/rapidfireai/dispatcher/__init__.py b/rapidfireai/dispatcher/__init__.py new file mode 100644 index 00000000..415dda1c --- /dev/null +++ b/rapidfireai/dispatcher/__init__.py @@ -0,0 +1,9 @@ +""" +RapidFire AI Dispatcher + +REST API for Interactive Control of experiments. +""" + +from rapidfireai.dispatcher.dispatcher import Dispatcher, run_dispatcher, serve_forever, start_dispatcher_thread + +__all__ = ["Dispatcher", "run_dispatcher", "serve_forever", "start_dispatcher_thread"] diff --git a/rapidfireai/dispatcher/dispatcher.py b/rapidfireai/dispatcher/dispatcher.py new file mode 100644 index 00000000..5a73fc51 --- /dev/null +++ b/rapidfireai/dispatcher/dispatcher.py @@ -0,0 +1,740 @@ +""" +RapidFire AI Dispatcher + +REST API for Interactive Control of experiments. +Provides HTTP endpoints for dynamic run/pipeline management during experiment execution. +""" + +import json +import logging +import os +import threading +import traceback + +from flask import Flask, Response, jsonify, request +from flask_cors import CORS + +from rapidfireai.db.rf_db import RfDb +from rapidfireai.utils.constants import ( + ColabConfig, + DispatcherConfig, + FrontendConfig, + ICOperation, + MLflowConfig, + RF_LOG_FILENAME, + RF_LOG_PATH, +) +from rapidfireai.utils.dispatcher_utils import check_experiment_running + +CORS_ALLOWED_ORIGINS = ["http://localhost", DispatcherConfig.URL, MLflowConfig.URL, FrontendConfig.URL, "*"] + + +class Dispatcher: + """ + REST API server for interactive control of experiments. + Provides endpoints for: + - Run management: /dispatcher/get-all-runs, /dispatcher/stop-run, etc. + - Pipeline management: /dispatcher/get-all-pipelines, /dispatcher/stop-pipeline, etc. + """ + + # Status mapping for frontend compatibility (lowercase to capitalized) + STATUS_MAP = { + "new": "New", + "ongoing": "Ongoing", + "completed": "Completed", + "stopped": "Stopped", + "deleted": "Deleted", + "failed": "Failed", + } + + def __init__(self): + """Initialize the Dispatcher with database connection and Flask app.""" + self.db = RfDb() + self.app = Flask(__name__) + + # CORS configuration + cors_credentials = True if ColabConfig.ON_COLAB else False + CORS( + self.app, + resources={ + r"/*": { + "origins": "*", + "methods": ["GET", "POST", "PUT", "DELETE", "OPTIONS"], + "allow_headers": ["Content-Type", "Authorization"], + "expose_headers": ["Content-Type"], + "supports_credentials": cors_credentials, + } + }, + ) + + self.register_routes() + + def register_routes(self): + """Register all REST API routes.""" + route_prefix = "/dispatcher" + + # Handle OPTIONS preflight globally + @self.app.before_request + def handle_preflight(): + if request.method == "OPTIONS": + response = jsonify({}) + response.headers.add("Access-Control-Allow-Origin", "*") + response.headers.add("Access-Control-Allow-Headers", "Content-Type,Authorization") + response.headers.add("Access-Control-Allow-Methods", "GET,PUT,POST,DELETE,OPTIONS") + response.headers.add("Access-Control-Max-Age", "3600") + return response + + # Health check + self.app.add_url_rule(f"{route_prefix}/health-check", "health_check", self.health_check, methods=["GET"]) + + # ============================================================ + # SHARED ROUTES (both modes) + # ============================================================ + self.app.add_url_rule( + f"{route_prefix}/get-running-experiment", + "get_running_experiment", + self.get_running_experiment, + methods=["GET"], + ) + 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-experiment-logs", + "get_experiment_logs", + self.get_experiment_logs, + methods=["POST"], + ) + self.app.add_url_rule( + f"{route_prefix}/get-ic-logs", + "get_ic_logs", + self.get_ic_logs, + methods=["POST"], + ) + self.app.add_url_rule( + f"{route_prefix}/is-experiment-running", + "is_experiment_running", + self.is_experiment_running, + methods=["POST"], + ) + + # ============================================================ + # FIT MODE ROUTES (runs) + # ============================================================ + 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}/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}/clone-modify-run", + "clone_modify_run", + self.clone_modify_run, + methods=["POST"], + ) + + # ============================================================ + # EVALS MODE ROUTES (pipelines) + # ============================================================ + self.app.add_url_rule( + f"{route_prefix}/get-all-pipelines", + "get_all_pipelines", + self.get_all_pipelines, + methods=["GET"], + ) + self.app.add_url_rule( + f"{route_prefix}/list-all-pipeline-ids", + "list_all_pipeline_ids", + self.list_all_pipeline_ids, + methods=["GET"], + ) + self.app.add_url_rule(f"{route_prefix}/get-pipeline", "get_pipeline", self.get_pipeline, methods=["POST"]) + self.app.add_url_rule( + f"{route_prefix}/get-pipeline-config-json/", + "get_pipeline_config_json", + self.get_pipeline_config_json, + methods=["GET"], + ) + self.app.add_url_rule( + f"{route_prefix}/stop-pipeline", + "stop_pipeline", + self.stop_pipeline, + methods=["POST"], + ) + self.app.add_url_rule( + f"{route_prefix}/resume-pipeline", + "resume_pipeline", + self.resume_pipeline, + methods=["POST"], + ) + self.app.add_url_rule( + f"{route_prefix}/delete-pipeline", + "delete_pipeline", + self.delete_pipeline, + methods=["POST"], + ) + self.app.add_url_rule( + f"{route_prefix}/clone-pipeline", + "clone_pipeline", + self.clone_pipeline, + methods=["POST"], + ) + + # IC operations status + self.app.add_url_rule( + f"{route_prefix}/operation-status/", + "get_operation_status", + self.get_operation_status, + methods=["GET"], + ) + self.app.add_url_rule( + f"{route_prefix}/all-operations", + "get_all_operations", + self.get_all_operations, + methods=["GET"], + ) + + # ============================================================ + # HEALTH CHECK + # ============================================================ + + def health_check(self) -> tuple[Response, int]: + """Health check endpoint.""" + return jsonify({"status": "ok", "message": "Dispatcher is running"}), 200 + + # ============================================================ + # SHARED ROUTES + # ============================================================ + + def get_running_experiment(self) -> tuple[Response, int]: + """Get the currently running experiment.""" + try: + experiment = self.db.get_running_experiment() + + # Detect mode by checking if we have runs or pipelines + has_runs = len(self.db.get_all_runs()) > 0 + has_pipelines = len(self.db.get_all_pipelines()) > 0 + mode = "fit" if has_runs else ("evals" if has_pipelines else "unknown") + + return jsonify({ + "experiment_id": experiment["experiment_id"], + "experiment_name": experiment["experiment_name"], + "status": experiment["status"].value if hasattr(experiment["status"], "value") else experiment["status"], + "metric_experiment_id": experiment.get("metric_experiment_id"), + "mode": mode, + }), 200 + except Exception as e: + return jsonify({"error": str(e)}), 500 + + def get_all_experiment_names(self) -> tuple[Response, int]: + """Get all experiment names.""" + try: + names = self.db.get_all_experiment_names() + return jsonify(names), 200 + except Exception as e: + return jsonify({"error": str(e)}), 500 + + def is_experiment_running(self) -> tuple[Response, int]: + """Check if a specific experiment is currently running.""" + try: + data = request.get_json() + if not data or "experiment_name" not in data: + return jsonify({"error": "experiment_name is required"}), 400 + + is_running = check_experiment_running(self.db, data["experiment_name"]) + return jsonify({"is_running": is_running}), 200 + except Exception: + # If anything fails, assume experiment is not running (safer to disable button) + return jsonify({"is_running": False}), 200 + + def get_experiment_logs(self) -> tuple[Response, int]: + """Get experiment logs.""" + 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: + try: + running_exp = self.db.get_running_experiment() + experiment_name = running_exp["experiment_name"] + except Exception: + return jsonify([]), 200 + + log_file_path = os.path.join(RF_LOG_PATH, experiment_name, RF_LOG_FILENAME) + + if not os.path.exists(log_file_path): + return jsonify([]), 200 + + experiment_logs = [] + with open(log_file_path, encoding="utf-8") as f: + for line in f: + if f"[{experiment_name}:" in line or f"| {experiment_name} |" in line: + experiment_logs.append(line.strip()) + + return jsonify(experiment_logs), 200 + except Exception as e: + return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 + + def get_ic_logs(self) -> tuple[Response, int]: + """Get interactive control logs.""" + 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: + try: + running_exp = self.db.get_running_experiment() + experiment_name = running_exp["experiment_name"] + except Exception: + return jsonify([]), 200 + + log_file_path = os.path.join(RF_LOG_PATH, experiment_name, RF_LOG_FILENAME) + + if not os.path.exists(log_file_path): + return jsonify([]), 200 + + ic_logs = [] + with open(log_file_path, encoding="utf-8") as f: + for line in f: + if f"| {experiment_name} | interactive-control |" in line: + ic_logs.append(line.strip()) + + return jsonify(ic_logs), 200 + except Exception as e: + return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 + + # ============================================================ + # FIT MODE ROUTES (runs) + # ============================================================ + + def _format_run_for_response(self, run_id: int, run_data: dict) -> dict: + """Format run data for API response.""" + config_leaf = run_data.get("config_leaf", {}) + if config_leaf: + config_leaf = dict(config_leaf) + config_leaf.pop("additional_kwargs", None) + if "peft_params" in config_leaf: + config_leaf["peft_params"].pop("task_type", None) + if "model_kwargs" in config_leaf: + config_leaf["model_kwargs"].pop("torch_dtype", None) + if "reward_funcs" in config_leaf: + config_leaf.pop("reward_funcs", None) + + status = run_data["status"] + status_value = status.value if hasattr(status, "value") else status + + return { + "run_id": run_id, + "status": status_value, + "metric_run_id": run_data.get("metric_run_id"), + "config": config_leaf, + "flattened_config": run_data.get("flattened_config", {}), + "completed_steps": run_data.get("completed_steps", 0), + "total_steps": run_data.get("total_steps", 0), + "num_epochs_completed": run_data.get("num_epochs_completed", 0), + "error": run_data.get("error", ""), + "source": run_data["source"].value if run_data.get("source") else None, + "ended_by": run_data["ended_by"].value if run_data.get("ended_by") else None, + } + + def get_all_runs(self) -> tuple[Response, int]: + """Get all runs (fit mode).""" + try: + runs = self.db.get_all_runs() + result = [self._format_run_for_response(run_id, run_data) for run_id, run_data in runs.items()] + return jsonify(result), 200 + except Exception as e: + return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 + + def get_run(self) -> tuple[Response, int]: + """Get a specific run (fit mode).""" + try: + data = request.get_json() + run_id = data.get("run_id") + if not run_id: + return jsonify({"error": "run_id is required"}), 400 + + run_data = self.db.get_run(run_id) + return jsonify(self._format_run_for_response(run_id, run_data)), 200 + except Exception as e: + return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 + + def stop_run(self) -> tuple[Response, int]: + """Stop a run (fit mode).""" + try: + data = request.get_json() + run_id = data.get("run_id") + if not run_id: + return jsonify({"error": "run_id is required"}), 400 + + ic_id = self.db.create_ic_operation( + target_id=run_id, + target_type="run", + operation=ICOperation.STOP, + ) + return jsonify({"ic_id": ic_id}), 200 + except Exception as e: + return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 + + def resume_run(self) -> tuple[Response, int]: + """Resume a run (fit mode).""" + try: + data = request.get_json() + run_id = data.get("run_id") + if not run_id: + return jsonify({"error": "run_id is required"}), 400 + + ic_id = self.db.create_ic_operation( + target_id=run_id, + target_type="run", + operation=ICOperation.RESUME, + ) + return jsonify({"ic_id": ic_id}), 200 + except Exception as e: + return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 + + def delete_run(self) -> tuple[Response, int]: + """Delete a run (fit mode).""" + try: + data = request.get_json() + run_id = data.get("run_id") + if not run_id: + return jsonify({"error": "run_id is required"}), 400 + + ic_id = self.db.create_ic_operation( + target_id=run_id, + target_type="run", + operation=ICOperation.DELETE, + ) + return jsonify({"ic_id": ic_id}), 200 + except Exception as e: + return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 + + def clone_modify_run(self) -> tuple[Response, int]: + """Clone and modify a run (fit mode).""" + try: + data = request.get_json() + if not data: + return jsonify({"error": "No JSON data received"}), 400 + + run_id = data.get("run_id") + if not run_id: + return jsonify({"error": "run_id is required"}), 400 + + config = data.get("config") + if not config: + return jsonify({"error": "config is required"}), 400 + + warm_start = data.get("warm_start", False) + operation = ICOperation.CLONE_WARM if warm_start else ICOperation.CLONE + + ic_id = self.db.create_ic_operation( + target_id=run_id, + target_type="run", + operation=operation, + config_data=config, + ) + return jsonify({"ic_id": ic_id}), 200 + except Exception as e: + return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 + + # ============================================================ + # EVALS MODE ROUTES (pipelines) + # ============================================================ + + def _format_pipeline_for_response(self, pipeline: dict) -> dict: + """Format pipeline data for API response.""" + status = pipeline.get("status", "") + status_value = status.value if hasattr(status, "value") else status + + # Parse flattened_config to use as 'config' field for IC compatibility + flattened_config = pipeline.get("flattened_config", {}) + if isinstance(flattened_config, str): + try: + import json + flattened_config = json.loads(flattened_config) + except (json.JSONDecodeError, ValueError): + flattened_config = {} + + return { + "pipeline_id": pipeline.get("pipeline_id"), + "context_id": pipeline.get("context_id"), + "pipeline_type": pipeline.get("pipeline_type"), + "pipeline_config_json": pipeline.get("pipeline_config_json", {}), + "flattened_config": flattened_config, + "config": flattened_config, # Add 'config' field for InteractiveController compatibility + "status": status_value, + "current_shard_id": pipeline.get("current_shard_id", 0), + "shards_completed": pipeline.get("shards_completed", 0), + "total_samples_processed": pipeline.get("total_samples_processed", 0), + "metric_run_id": pipeline.get("metric_run_id"), + "error": pipeline.get("error", ""), + "created_at": pipeline.get("created_at"), + } + + def get_all_pipelines(self) -> tuple[Response, int]: + """Get all pipelines (evals mode).""" + try: + pipelines = self.db.get_all_pipelines() + result = [self._format_pipeline_for_response(p) for p in pipelines] + return jsonify(result), 200 + except Exception as e: + return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 + + def list_all_pipeline_ids(self) -> tuple[Response, int]: + """Get lightweight list of pipeline IDs (evals mode).""" + try: + pipelines = self.db.get_all_pipeline_ids() + return jsonify(pipelines), 200 + except Exception as e: + return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 + + def get_pipeline(self) -> tuple[Response, int]: + """Get a specific pipeline (evals mode).""" + try: + data = request.get_json() + pipeline_id = data.get("pipeline_id") + if not pipeline_id: + return jsonify({"error": "pipeline_id is required"}), 400 + + pipeline = self.db.get_pipeline(pipeline_id) + if not pipeline: + return jsonify({"error": f"Pipeline {pipeline_id} not found"}), 404 + + return jsonify(self._format_pipeline_for_response(pipeline)), 200 + except Exception as e: + return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 + + def get_pipeline_config_json(self, pipeline_id: int) -> tuple[Response, int]: + """Get config JSON for a pipeline (evals mode).""" + try: + config_data = self.db.get_pipeline_config_json(pipeline_id) + if not config_data: + return jsonify({"error": f"Pipeline {pipeline_id} not found"}), 404 + return jsonify(config_data), 200 + except Exception as e: + return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 + + def stop_pipeline(self) -> tuple[Response, int]: + """Stop a pipeline (evals mode).""" + try: + data = request.get_json() + pipeline_id = data.get("pipeline_id") + if not pipeline_id: + return jsonify({"error": "pipeline_id is required"}), 400 + + pipeline = self.db.get_pipeline(pipeline_id) + if not pipeline: + return jsonify({"error": f"Pipeline {pipeline_id} not found"}), 404 + + ic_id = self.db.create_ic_operation( + target_id=pipeline_id, + target_type="pipeline", + operation=ICOperation.STOP, + ) + return jsonify({"ic_id": ic_id, "message": f"Stop request created for pipeline {pipeline_id}"}), 200 + except Exception as e: + return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 + + def resume_pipeline(self) -> tuple[Response, int]: + """Resume a pipeline (evals mode).""" + try: + data = request.get_json() + pipeline_id = data.get("pipeline_id") + if not pipeline_id: + return jsonify({"error": "pipeline_id is required"}), 400 + + pipeline = self.db.get_pipeline(pipeline_id) + if not pipeline: + return jsonify({"error": f"Pipeline {pipeline_id} not found"}), 404 + + ic_id = self.db.create_ic_operation( + target_id=pipeline_id, + target_type="pipeline", + operation=ICOperation.RESUME, + ) + return jsonify({"ic_id": ic_id, "message": f"Resume request created for pipeline {pipeline_id}"}), 200 + except Exception as e: + return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 + + def delete_pipeline(self) -> tuple[Response, int]: + """Delete a pipeline (evals mode).""" + try: + data = request.get_json() + pipeline_id = data.get("pipeline_id") + if not pipeline_id: + return jsonify({"error": "pipeline_id is required"}), 400 + + pipeline = self.db.get_pipeline(pipeline_id) + if not pipeline: + return jsonify({"error": f"Pipeline {pipeline_id} not found"}), 404 + + ic_id = self.db.create_ic_operation( + target_id=pipeline_id, + target_type="pipeline", + operation=ICOperation.DELETE, + ) + return jsonify({"ic_id": ic_id, "message": f"Delete request created for pipeline {pipeline_id}"}), 200 + except Exception as e: + return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 + + def clone_pipeline(self) -> tuple[Response, int]: + """Clone a pipeline (evals mode).""" + try: + data = request.get_json() + if not data: + return jsonify({"error": "No JSON data provided"}), 400 + + parent_pipeline_id = data.get("parent_pipeline_id") + if not parent_pipeline_id: + return jsonify({"error": "parent_pipeline_id is required"}), 400 + + config_json = data.get("config_json") + if not config_json: + return jsonify({"error": "config_json is required"}), 400 + + parent_pipeline = self.db.get_pipeline(parent_pipeline_id) + if not parent_pipeline: + return jsonify({"error": f"Parent pipeline {parent_pipeline_id} not found"}), 404 + + request_data = { + "parent_pipeline_id": parent_pipeline_id, + "config_json": config_json, + } + + ic_id = self.db.create_ic_operation( + target_id=0, # Will be set when the new pipeline is created + target_type="pipeline", + operation=ICOperation.CLONE, + config_data=request_data, + ) + return jsonify({ + "ic_id": ic_id, + "message": f"Clone request created from parent pipeline {parent_pipeline_id}", + }), 200 + except Exception as e: + return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 + + # ============================================================ + # IC OPERATIONS STATUS + # ============================================================ + + def get_operation_status(self, ic_id: int) -> tuple[Response, int]: + """Get status of a specific IC operation.""" + try: + operation = self.db.get_ic_operation(ic_id) + if not operation: + return jsonify({"error": f"Operation {ic_id} not found"}), 404 + return jsonify(operation), 200 + except Exception as e: + return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 + + def get_all_operations(self) -> tuple[Response, int]: + """Get all IC operations.""" + try: + operations = self.db.get_all_ic_operations() + return jsonify(operations), 200 + except Exception as e: + return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 + + +# ============================================================================ +# SERVER FUNCTIONS +# ============================================================================ + +def run_dispatcher(host: str = "0.0.0.0", port: int = 8851) -> None: + """Run the dispatcher server (blocking).""" + try: + from waitress import serve + + dispatcher = Dispatcher() + logging.getLogger("waitress").setLevel(logging.WARNING) + serve(dispatcher.app, host=host, port=port, threads=6) + except Exception as e: + print(f"CRITICAL: Dispatcher crashed: {e}") + traceback.print_exc() + + +# Global dispatcher thread tracking +_dispatcher_thread: threading.Thread | None = None +_dispatcher_lock = threading.Lock() + + +def _check_port_in_use(host: str, port: int) -> bool: + """Check if a port is already in use.""" + import socket + + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + return s.connect_ex((host, port)) == 0 + + +def start_dispatcher_thread(host: str = "0.0.0.0", port: int = 8851, logger=None) -> threading.Thread | None: + """ + Start the dispatcher REST API server in a background daemon thread. + + Args: + host: Host to bind to + port: Port to bind to + logger: Optional logger instance + + Returns: + The dispatcher thread object, or None if startup failed + """ + global _dispatcher_thread + + with _dispatcher_lock: + # Check if our dispatcher thread is already running + if _dispatcher_thread is not None and _dispatcher_thread.is_alive(): + msg = f"Dispatcher thread already running on port {port}, reusing existing thread" + if logger: + logger.info(msg) + return _dispatcher_thread + + try: + if _check_port_in_use(host, port): + msg = f"Port {port} is already in use" + if logger: + logger.warning(msg) + return None + + _dispatcher_thread = threading.Thread( + target=run_dispatcher, + kwargs={"host": host, "port": port}, + daemon=True, + name="DispatcherThread", + ) + _dispatcher_thread.start() + + msg = f"Started interactive control dispatcher on http://{host}:{port}" + if logger: + logger.info(msg) + else: + print(msg) + + return _dispatcher_thread + + except Exception as e: + error_msg = f"Failed to start dispatcher: {e}" + if logger: + logger.warning(error_msg) + else: + print(f"WARNING: {error_msg}") + return None + + +def serve_forever() -> Flask: + """Start the Dispatcher via Gunicorn (WSGI entry point).""" + return Dispatcher().app + + +if __name__ == "__main__": + run_dispatcher() diff --git a/rapidfireai/fit/dispatcher/gunicorn.conf.py b/rapidfireai/dispatcher/gunicorn.conf.py similarity index 87% rename from rapidfireai/fit/dispatcher/gunicorn.conf.py rename to rapidfireai/dispatcher/gunicorn.conf.py index 5f7241e2..e31f6318 100644 --- a/rapidfireai/fit/dispatcher/gunicorn.conf.py +++ b/rapidfireai/dispatcher/gunicorn.conf.py @@ -1,13 +1,13 @@ """Gunicorn configuration file for the RapidFire dispatcher""" -from rapidfireai.fit.db.rf_db import RfDb +from rapidfireai.db import RfDb from rapidfireai.utils.constants import DispatcherConfig # Other Gunicorn settings... bind = f"{DispatcherConfig.HOST}:{DispatcherConfig.PORT}" workers = 1 # Single worker for Colab/single-user environments to save memory -wsgi_app = "rapidfireai.fit.dispatcher.dispatcher:serve_forever()" +wsgi_app = "rapidfireai.dispatcher.dispatcher:serve_forever()" def on_starting(server): diff --git a/rapidfireai/evals/actors/doc_actor.py b/rapidfireai/evals/actors/doc_actor.py index 16ad017b..11ca9a4e 100644 --- a/rapidfireai/evals/actors/doc_actor.py +++ b/rapidfireai/evals/actors/doc_actor.py @@ -18,7 +18,7 @@ from rapidfireai.evals.rag.prompt_manager import PromptManager from rapidfireai.evals.rag.rag_pipeline import LangChainRagSpec -from rapidfireai.evals.utils.logger import RFLogger +from rapidfireai.utils.logging import RFLogger @ray.remote diff --git a/rapidfireai/evals/actors/query_actor.py b/rapidfireai/evals/actors/query_actor.py index d6ec4fa1..2cbd72ae 100644 --- a/rapidfireai/evals/actors/query_actor.py +++ b/rapidfireai/evals/actors/query_actor.py @@ -14,18 +14,15 @@ import ray import pickle -from rapidfireai.utils.constants import MLflowConfig -from rapidfireai.evals.utils.constants import SEARCH_DEFAULTS, VALID_SEARCH_TYPES +from rapidfireai.utils.constants import MLflowConfig, SEARCH_DEFAULTS, VALID_SEARCH_TYPES, RF_EXPERIMENT_PATH from langchain_community.vectorstores import FAISS from langchain_postgres import PGVector from langchain_pinecone import PineconeVectorStore from pinecone import Pinecone -from rapidfireai.utils.constants import RF_EXPERIMENT_PATH from rapidfireai.evals.actors.inference_engines import InferenceEngine from rapidfireai.evals.rag.rag_pipeline import LangChainRagSpec -from rapidfireai.evals.rag.prompt_manager import PromptManager -from rapidfireai.evals.utils.logger import RFLogger +from rapidfireai.utils.logging import RFLogger @ray.remote @@ -67,9 +64,10 @@ def __init__( if "CUDA_VISIBLE_DEVICES" in os.environ and os.environ["CUDA_VISIBLE_DEVICES"]: try: import torch + if torch.cuda.is_available(): # Force CUDA initialization by performing a simple operation - _ = torch.zeros(1, device='cuda') + _ = torch.zeros(1, device="cuda") torch.cuda.synchronize() except Exception: # Silently continue if CUDA initialization fails (will use CPU) diff --git a/rapidfireai/evals/actors/rate_limiter_actor.py b/rapidfireai/evals/actors/rate_limiter_actor.py index 38dd0049..ea219c5d 100644 --- a/rapidfireai/evals/actors/rate_limiter_actor.py +++ b/rapidfireai/evals/actors/rate_limiter_actor.py @@ -4,11 +4,12 @@ Provides a single point of coordination for rate limiting across all distributed query processing actors. """ + import ray -from rapidfireai.utils.constants import RF_EXPERIMENT_PATH +from rapidfireai.utils.logging import RFLogger from rapidfireai.evals.utils.ratelimiter import OpenAIRateLimiter, RequestStatus -from rapidfireai.evals.utils.logger import RFLogger +from rapidfireai.utils.constants import RF_EXPERIMENT_PATH @ray.remote @@ -102,4 +103,3 @@ async def get_stats(self) -> dict: # Export for use in other modules __all__ = ["RateLimiterActor"] - diff --git a/rapidfireai/evals/db/README.md b/rapidfireai/evals/db/README.md deleted file mode 100644 index 6feb13ba..00000000 --- a/rapidfireai/evals/db/README.md +++ /dev/null @@ -1,131 +0,0 @@ -# RF-Inferno Database Module - -Provides database interface for tracking multi-pipeline inference experiments. - -## Files - -- `tables.sql` - SQLite schema definition -- `db_interface.py` - Low-level database connection and query execution -- `rf_db.py` - High-level CRUD operations for experiments, contexts, pipelines, and tasks -- `__init__.py` - Module exports - -## Quick Start - -```python -from rapidfireai.evals.db import RFDatabase - -# Initialize database -db = RFDatabase() - -# Create experiment -exp_id = db.create_experiment( - experiment_name="my_experiment", - num_shards=5, - num_actors=4, - num_cpus=8, - num_gpus=2 -) - -# Create pipeline -pipeline_id = db.create_pipeline( - experiment_id=exp_id, - pipeline_name="baseline_temp0.7", - pipeline_type="vllm", - model_config_json='{"model": "Qwen/Qwen2.5-3B-Instruct"}', - sampling_params_json='{"temperature": 0.7}', - status="new" -) - -# Update progress -db.set_pipeline_progress( - pipeline_id, - current_shard_id=2, - shards_completed=2, - total_samples_processed=1024 -) - -# Create task -task_id = db.create_actor_task( - experiment_id=exp_id, - pipeline_id=pipeline_id, - actor_id=0, - shard_id=2 -) - -# Close connection -db.close() -``` - -## API Methods - -### Experiments -- `create_experiment(name, num_shards, num_actors, ...) -> exp_id` -- `get_running_experiment() -> dict` -- `get_all_experiment_names() -> list[str]` -- `set_experiment_status(exp_id, status)` -- `set_experiment_error(exp_id, error)` -- `get_experiment_error(exp_id) -> str` -- `reset_all_tables(experiments_table=False)` - -### Contexts (RAG) -- `create_context(hash, rag_json, prompt_json) -> context_id` -- `get_context(context_id) -> dict` -- `set_context_status(context_id, status)` -- `set_context_start_time(context_id, time)` -- `set_context_end_time(context_id, time, duration)` -- `set_context_error(context_id, error)` - -### Pipelines -- `create_pipeline(exp_id, name, type, model_json, sampling_json) -> pipeline_id` -- `get_pipeline(pipeline_id) -> dict` -- `get_all_pipelines(exp_id) -> list[dict]` -- `set_pipeline_status(pipeline_id, status)` -- `set_pipeline_progress(pipeline_id, shard_id, completed, samples)` -- `set_pipeline_error(pipeline_id, error)` - -### Actor Tasks -- `create_actor_task(exp_id, pipeline_id, actor_id, shard_id) -> task_id` -- `get_actor_task(task_id) -> dict` -- `get_running_actor_tasks(exp_id) -> list[dict]` -- `set_actor_task_status(task_id, status)` -- `set_actor_task_start_time(task_id, time)` -- `set_actor_task_end_time(task_id, time, duration)` -- `set_actor_task_error(task_id, error)` - -## Status Values - -### Experiments -- `running` - Currently executing -- `completed` - Successfully finished -- `failed` - Encountered error -- `cancelled` - User cancelled - -### Contexts -- `pending` - Not yet built -- `building` - Index construction in progress -- `ready` - Ready for use -- `failed` - Build failed - -### Pipelines -- `new` - Just created -- `ongoing` - Currently processing shards -- `completed` - All shards processed -- `stopped` - User stopped -- `deleted` - User deleted -- `failed` - Encountered error - -### Actor Tasks -- `scheduled` - Task assigned but not started -- `in_progress` - Currently executing -- `completed` - Successfully finished -- `failed` - Encountered error - -## Testing - -Run the test suite: -```bash -python tests/test_rf_db.py -``` - - - diff --git a/rapidfireai/evals/db/__init__.py b/rapidfireai/evals/db/__init__.py deleted file mode 100644 index 99b7e31f..00000000 --- a/rapidfireai/evals/db/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""RF-Inferno database module.""" - -from rapidfireai.evals.db.db_interface import DatabaseInterface -from rapidfireai.evals.db.rf_db import RFDatabase - -__all__ = ["DatabaseInterface", "RFDatabase"] diff --git a/rapidfireai/evals/db/rf_db.py b/rapidfireai/evals/db/rf_db.py deleted file mode 100644 index 0439fc80..00000000 --- a/rapidfireai/evals/db/rf_db.py +++ /dev/null @@ -1,1146 +0,0 @@ -""" -RF-Inferno Database Manager. - -Provides high-level interface for CRUD operations on the experiment database. -""" - -import json -import os -from typing import Any - -from rapidfireai.evals.db.db_interface import DatabaseInterface -from rapidfireai.evals.utils.constants import ( - ContextStatus, - ExperimentStatus, - PipelineStatus, - TaskStatus, -) -from rapidfireai.evals.utils.serialize import ( - decode_db_payload, - encode_payload, - extract_pipeline_config_json, -) - - -class RFDatabase: - """Database manager for RF-Inferno experiments.""" - - def __init__(self, db_interface: DatabaseInterface = None): - """ - Initialize the database manager. - - Args: - db_interface: Optional DatabaseInterface instance. If not provided, creates a new one. - """ - self.db = db_interface if db_interface else DatabaseInterface() - self._initialize_schema() - - def _initialize_schema(self): - """Initialize database schema from tables.sql file.""" - schema_path = os.path.join(os.path.dirname(__file__), "tables.sql") - if os.path.exists(schema_path): - with open(schema_path) as f: - schema_sql = f.read() - self.db.conn.executescript(schema_sql) - self.db.conn.commit() - - # Migration: Add metric_run_id to pipelines table if they don't exist - try: - cursor = self.db.conn.execute("PRAGMA table_info(pipelines)") - columns = [row[1] for row in cursor.fetchall()] - if "metric_run_id" not in columns: - self.db.conn.execute("ALTER TABLE pipelines ADD COLUMN metric_run_id TEXT") - self.db.conn.commit() - except Exception: - pass - - # Migration: Add metric_experiment_id to experiments table if they don't exist - try: - cursor = self.db.conn.execute("PRAGMA table_info(experiments)") - columns = [row[1] for row in cursor.fetchall()] - if "metric_experiment_id" not in columns: - self.db.conn.execute("ALTER TABLE experiments ADD COLUMN metric_experiment_id TEXT") - self.db.conn.commit() - except Exception: - pass - - def close(self): - """Close the database connection.""" - self.db.close() - - def create_tables(self): - """Create database tables. - - Public method for Gunicorn on_starting() callback. - Re-runs _initialize_schema() which uses CREATE TABLE IF NOT EXISTS, - making it safe to call multiple times. - """ - self._initialize_schema() - - # ============================================================================ - # EXPERIMENTS TABLE METHODS - # ============================================================================ - - def create_experiment( - self, - experiment_name: str, - num_actors: int, - num_cpus: int = None, - num_gpus: int = None, - metric_experiment_id: str = None, - status: ExperimentStatus = ExperimentStatus.RUNNING, - num_shards: int = 0, - ) -> int: - """ - Create a new experiment record. - - Args: - experiment_name: Name of the experiment - num_actors: Number of query processing actors - num_cpus: Number of CPUs allocated - num_gpus: Number of GPUs allocated - metric_experiment_id: Optional MetricLogger experiment ID - status: Initial status (default: ExperimentStatus.RUNNING) - num_shards: Number of shards for the dataset (default: 0) - - Returns: - experiment_id of the created experiment - """ - query = """ - INSERT INTO experiments ( - experiment_name, num_actors, num_shards, num_cpus, num_gpus, - metric_experiment_id, status, error - ) VALUES (?, ?, ?, ?, ?, ?, ?, '') - """ - self.db.execute( - query, - ( - experiment_name, - num_actors, - num_shards, - num_cpus, - num_gpus, - metric_experiment_id, - status.value, - ), - commit=True, - ) - return self.db.cursor.lastrowid - - def set_experiment_status(self, experiment_id: int, status: ExperimentStatus): - """ - Update experiment status. - - Args: - experiment_id: ID of the experiment - status: New status (ExperimentStatus enum) - """ - query = "UPDATE experiments SET status = ? WHERE experiment_id = ?" - self.db.execute(query, (status.value, experiment_id), commit=True) - - def set_experiment_error(self, experiment_id: int, error: str): - """ - Set error message for an experiment. - - Args: - experiment_id: ID of the experiment - error: Error message - """ - query = "UPDATE experiments SET error = ? WHERE experiment_id = ?" - self.db.execute(query, (error, experiment_id), commit=True) - - def set_experiment_num_shards(self, experiment_id: int, num_shards: int): - """ - Update the number of shards for an experiment. - - Args: - experiment_id: ID of the experiment - num_shards: Number of shards - """ - query = "UPDATE experiments SET num_shards = ? WHERE experiment_id = ?" - self.db.execute(query, (num_shards, experiment_id), commit=True) - - def set_experiment_resources( - self, - experiment_id: int, - num_actors: int, - num_cpus: int = None, - num_gpus: int = None, - ): - """ - Update resource allocation for an experiment. - - Args: - experiment_id: ID of the experiment - num_actors: Number of actors - num_cpus: Number of CPUs (optional) - num_gpus: Number of GPUs (optional) - """ - query = "UPDATE experiments SET num_actors = ?, num_cpus = ?, num_gpus = ? WHERE experiment_id = ?" - self.db.execute( - query, (num_actors, num_cpus, num_gpus, experiment_id), commit=True - ) - - def reset_all_tables(self, experiments_table: bool = False) -> None: - """ - Clear data from experiment tables. - - Args: - experiments_table: If True, also clear the experiments table (default: False) - """ - # Clear dependent tables first (due to foreign keys) - tables = ["actor_tasks", "contexts", "interactive_control", "pipelines"] - - for table in tables: - self.db.execute(f"DELETE FROM {table}", commit=True) - - # Optionally clear experiments table - if experiments_table: - self.db.execute("DELETE FROM experiments", commit=True) - tables.append("experiments") - - # Reset auto-increment indices - for table in tables: - self.db.execute("DELETE FROM sqlite_sequence WHERE name = ?", (table,), commit=True) - - def reset_experiment_states(self) -> None: - """ - Reset the experiment states when a running task is cancelled. - Marks ongoing/new pipelines as FAILED and their contexts as FAILED. - Similar to fit mode's reset_experiment_states(). - """ - from rapidfireai.evals.utils.constants import ContextStatus, PipelineStatus, TaskStatus - - # Mark all scheduled and in-progress actor tasks as failed - query = """ - UPDATE actor_tasks - 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 pipelines as failed - query = """ - UPDATE pipelines - SET status = ? - WHERE status = ? OR status = ? - """ - self.db.execute( - query, (PipelineStatus.FAILED.value, PipelineStatus.ONGOING.value, PipelineStatus.NEW.value), commit=True - ) - - # Mark ongoing and new contexts as failed - query = """ - UPDATE contexts - SET status = ? - WHERE status = ? OR status = ? - """ - self.db.execute( - query, (ContextStatus.FAILED.value, ContextStatus.ONGOING.value, ContextStatus.NEW.value), commit=True - ) - - # Reset all pending interactive control tasks - query = """ - UPDATE interactive_control - SET status = ? - WHERE status = ? - """ - self.db.execute(query, (TaskStatus.FAILED.value, TaskStatus.SCHEDULED.value), commit=True) - - def get_experiment(self, experiment_id: int) -> dict[str, Any] | None: - """ - Get experiment details by ID. - - Args: - experiment_id: ID of the experiment - - Returns: - Experiment dictionary with all fields, or None if not found - """ - query = """ - SELECT experiment_id, experiment_name, num_actors, num_cpus, num_gpus, - metric_experiment_id, status, num_shards, error, created_at - FROM experiments - WHERE experiment_id = ? - """ - result = self.db.execute(query, params=(experiment_id,), fetch=True) - if result and len(result) > 0: - row = result[0] - return { - "experiment_id": row[0], - "experiment_name": row[1], - "num_actors": row[2], - "num_cpus": row[3], - "num_gpus": row[4], - "metric_experiment_id": row[5], - "status": row[6], - "num_shards": row[7], - "error": row[8], - "created_at": row[9], - } - return None - - def get_experiment_error(self, experiment_id: int) -> str: - """ - Get error message for an experiment. - - Args: - experiment_id: ID of the experiment - - Returns: - Error message string - """ - query = "SELECT error FROM experiments WHERE experiment_id = ?" - result = self.db.execute(query, (experiment_id,), fetch=True) - return result[0][0] if result else "" - - def get_all_experiment_names(self) -> list[str]: - """ - Get all experiment names. - - Returns: - List of experiment names - """ - query = "SELECT experiment_name FROM experiments" - result = self.db.execute(query, fetch=True) - return [row[0] for row in result] if result else [] - - def get_running_experiment(self) -> dict[str, Any] | None: - """ - Get the currently running experiment (most recent if multiple). - - Returns: - Dictionary with all experiment fields, or None if no running experiment - """ - query = """ - SELECT experiment_id, experiment_name, metric_experiment_id, num_shards, - num_actors, num_cpus, num_gpus, status, error, created_at - FROM experiments - WHERE status = ? - ORDER BY experiment_id DESC - LIMIT 1 - """ - result = self.db.execute(query, (ExperimentStatus.RUNNING.value,), fetch=True) - if result: - row = result[0] - return { - "experiment_id": row[0], - "experiment_name": row[1], - "metric_experiment_id": row[2], - "num_shards": row[3], - "num_actors": row[4], - "num_cpus": row[5], - "num_gpus": row[6], - "status": row[7], - "error": row[8], - "created_at": row[9], - } - return None - - # ============================================================================ - # CONTEXTS TABLE METHODS - # ============================================================================ - - def create_context( - self, - context_hash: str, - rag_config_json: str = None, - prompt_config_json: str = None, - status: ContextStatus = ContextStatus.NEW, - ) -> int: - """ - Create a new context record. - - Args: - context_hash: SHA256 hash of the context configuration - rag_config_json: JSON string of RAG configuration - prompt_config_json: JSON string of prompt manager configuration - status: Initial status (default: ContextStatus.NEW) - - Returns: - context_id of the created context (or existing if hash matches) - """ - # Check if context with this hash already exists - query = "SELECT context_id FROM contexts WHERE context_hash = ?" - result = self.db.execute(query, (context_hash,), fetch=True) - if result: - return result[0][0] - - # Create new context - query = """ - INSERT INTO contexts (context_hash, rag_config_json, prompt_config_json, status, error) - VALUES (?, ?, ?, ?, '') - """ - self.db.execute( - query, - (context_hash, rag_config_json, prompt_config_json, status.value), - commit=True, - ) - return self.db.cursor.lastrowid - - def get_context(self, context_id: int) -> dict[str, Any] | None: - """ - Get context by ID. - - Args: - context_id: ID of the context - - Returns: - Dictionary with all context fields, or None if not found - """ - query = """ - SELECT context_id, context_hash, rag_config_json, prompt_config_json, - status, error, started_at, completed_at, duration_seconds - FROM contexts - WHERE context_id = ? - """ - result = self.db.execute(query, (context_id,), fetch=True) - if result: - row = result[0] - return { - "context_id": row[0], - "context_hash": row[1], - "rag_config_json": row[2], - "prompt_config_json": row[3], - "status": row[4], - "error": row[5], - "started_at": row[6], - "completed_at": row[7], - "duration_seconds": row[8], - } - return None - - def get_context_by_hash(self, context_hash: str) -> dict[str, Any] | None: - """ - Get context by hash. - - Args: - context_hash: SHA256 hash of the context configuration - - Returns: - Dictionary with all context fields, or None if not found - """ - query = """ - SELECT context_id, context_hash, rag_config_json, prompt_config_json, - status, error, started_at, completed_at, duration_seconds - FROM contexts - WHERE context_hash = ? - """ - result = self.db.execute(query, (context_hash,), fetch=True) - if result: - row = result[0] - return { - "context_id": row[0], - "context_hash": row[1], - "rag_config_json": row[2], - "prompt_config_json": row[3], - "status": row[4], - "error": row[5], - "started_at": row[6], - "completed_at": row[7], - "duration_seconds": row[8], - } - return None - - def set_context_status(self, context_id: int, status: ContextStatus): - """ - Update context status. - - Args: - context_id: ID of the context - status: New status (ContextStatus enum) - """ - query = "UPDATE contexts SET status = ? WHERE context_id = ?" - self.db.execute(query, (status.value, context_id), commit=True) - - def set_context_start_time(self, context_id: int, start_time: str): - """ - Set start time for context building. - - Args: - context_id: ID of the context - start_time: Start timestamp (use datetime.now().isoformat()) - """ - query = "UPDATE contexts SET started_at = ? WHERE context_id = ?" - self.db.execute(query, (start_time, context_id), commit=True) - - def set_context_end_time( - self, context_id: int, end_time: str, duration_seconds: float - ): - """ - Set end time and duration for context building. - - Args: - context_id: ID of the context - end_time: End timestamp (use datetime.now().isoformat()) - duration_seconds: Duration in seconds - """ - query = "UPDATE contexts SET completed_at = ?, duration_seconds = ? WHERE context_id = ?" - self.db.execute(query, (end_time, duration_seconds, context_id), commit=True) - - def set_context_error(self, context_id: int, error: str): - """ - Set error message for a context. - - Args: - context_id: ID of the context - error: Error message - """ - query = "UPDATE contexts SET error = ? WHERE context_id = ?" - self.db.execute(query, (error, context_id), commit=True) - - # ============================================================================ - # PIPELINES TABLE METHODS - # ============================================================================ - - def create_pipeline( - self, - pipeline_type: str, - pipeline_config: Any, - context_id: int = None, - status: PipelineStatus = PipelineStatus.NEW, - flattened_config: dict[str, Any] = None, - ) -> int: - """ - Create a new pipeline record. - - Args: - pipeline_name: Name/identifier for the pipeline - pipeline_type: Type of pipeline ('vllm', 'openai_api', etc.) - pipeline_config: Pipeline configuration object (with classes/functions - will be encoded for pipeline_config column, - and JSON-serialized for pipeline_config_json column) - context_id: Optional context ID for RAG - status: Initial status (default: PipelineStatus.NEW) - flattened_config: Flattened configuration dict for IC Ops panel display - - Returns: - pipeline_id of the created pipeline - """ - # Serialize the full pipeline config using encode_payload (includes functions/classes) - encoded_config = encode_payload(pipeline_config) - - # Extract JSON-serializable data (excludes functions/classes) - import json - - json_config_dict = extract_pipeline_config_json(pipeline_config) - json_config_str = json.dumps(json_config_dict) if json_config_dict is not None else "{}" - flattened_config_str = json.dumps(flattened_config) if flattened_config is not None else "{}" - - query = """ - INSERT INTO pipelines ( - context_id, pipeline_type, - pipeline_config, pipeline_config_json, flattened_config, status, error, - current_shard_id, shards_completed, total_samples_processed, metric_run_id - ) VALUES (?, ?, ?, ?, ?, ?, '', '', 0, 0, NULL) - """ - self.db.execute( - query, - ( - context_id, - pipeline_type, - encoded_config, - json_config_str, - flattened_config_str, - status.value, - ), - commit=True, - ) - return self.db.cursor.lastrowid - - def set_pipeline_progress(self, pipeline_id: int) -> dict[str, Any] | None: - """ - Get pipeline by ID (legacy method name - actually gets pipeline, not sets progress). - - Args: - pipeline_id: ID of the pipeline - - Returns: - Dictionary with all pipeline fields, or None if not found - """ - query = """ - SELECT pipeline_id, context_id, pipeline_type, - pipeline_config, pipeline_config_json, status, current_shard_id, - shards_completed, total_samples_processed, metric_run_id, error, created_at - FROM pipelines - WHERE pipeline_id = ? - """ - result = self.db.execute(query, (pipeline_id,), fetch=True) - if result: - row = result[0] - # Decode the pipeline config from the database (use pipeline_config column) - decoded_config = decode_db_payload(row[3]) if row[3] else None - # Parse JSON config for display/analytics - import json - - json_config = json.loads(row[4]) if row[4] else None - return { - "pipeline_id": row[0], - "context_id": row[1], - "pipeline_type": row[2], - "pipeline_config": decoded_config, # Use decoded config for actual pipeline object - "pipeline_config_json": json_config, # JSON version for display/analytics - "status": row[5], - "current_shard_id": row[6], - "shards_completed": row[7], - "total_samples_processed": row[8], - "metric_run_id": row[9], - "error": row[10], - "created_at": row[11], - } - return None - - def get_pipeline_by_metric_run_id(self, metric_run_id: str) -> dict[str, Any] | None: - """ - Get pipeline by its metric_run_id (MLflow/Trackio run UUID). - - Args: - metric_run_id: The metric tracking run ID (UUID string) - - Returns: - Pipeline dictionary, or None if not found - """ - query = """ - SELECT pipeline_id, context_id, pipeline_type, - pipeline_config, pipeline_config_json, flattened_config, status, current_shard_id, - shards_completed, total_samples_processed, metric_run_id, error, created_at - FROM pipelines - WHERE metric_run_id = ? - """ - result = self.db.execute(query, params=(metric_run_id,), fetch=True) - if result and len(result) > 0: - row = result[0] - decoded_config = decode_db_payload(row[3]) if row[3] else None - json_config = json.loads(row[4]) if row[4] else None - flattened_config = json.loads(row[5]) if row[5] else {} - return { - "pipeline_id": row[0], - "context_id": row[1], - "pipeline_type": row[2], - "pipeline_config": decoded_config, - "pipeline_config_json": json_config, - "flattened_config": flattened_config, - "status": row[6], - "current_shard_id": row[7], - "shards_completed": row[8], - "total_samples_processed": row[9], - "metric_run_id": row[10], - "error": row[11], - "created_at": row[12], - } - return None - - def get_pipeline(self, pipeline_id: int | str) -> dict[str, Any] | None: - """ - Get a single pipeline by ID. - - Args: - pipeline_id: ID of the pipeline to retrieve - - Returns: - Pipeline dictionary, or None if not found - """ - query = """ - SELECT pipeline_id, context_id, pipeline_type, - pipeline_config, pipeline_config_json, flattened_config,status, current_shard_id, - shards_completed, total_samples_processed, metric_run_id, error, created_at - FROM pipelines - WHERE pipeline_id = ? - """ - pipeline = None - if isinstance(pipeline_id, str): - # Try as MLflow run ID (UUID string) - pipeline = self.get_pipeline_by_metric_run_id(pipeline_id) - # Fallback: try parsing as int - if pipeline: - return pipeline - result = self.db.execute(query, params=(pipeline_id,), fetch=True) - if result and len(result) > 0: - row = result[0] - # Decode the pipeline config from the database (use pipeline_config column) - decoded_config = decode_db_payload(row[3]) if row[3] else None - # Parse JSON config for display/analytics - - json_config = json.loads(row[4]) if row[4] else None - flattened_config = json.loads(row[5]) if row[5] else {} - return { - "pipeline_id": row[0], - "context_id": row[1], - "pipeline_type": row[2], - "pipeline_config": decoded_config, # Use decoded config for actual pipeline object - "pipeline_config_json": json_config, # JSON version for display/analytics - "flattened_config": flattened_config, # Flattened configuration for IC Ops panel display - "status": row[6], - "current_shard_id": row[7], - "shards_completed": row[8], - "total_samples_processed": row[9], - "metric_run_id": row[10], - "error": row[11], - "created_at": row[12], - } - return None - - def get_all_pipeline_ids(self) -> list[dict[str, Any]]: - """ - Get lightweight list of all pipelines with minimal info (no config). - - Optimized for auto-polling - returns only IDs and status without deserializing configs. - - Returns: - List of dicts with: pipeline_id, status, shards_completed, total_samples_processed - """ - query = """ - SELECT pipeline_id, status, shards_completed, total_samples_processed - FROM pipelines - ORDER BY pipeline_id DESC - """ - result = self.db.execute(query, fetch=True) - pipelines = [] - if result: - for row in result: - pipelines.append( - { - "pipeline_id": row[0], - "status": row[1], - "shards_completed": row[2], - "total_samples_processed": row[3], - } - ) - return pipelines - - def get_pipeline_config_json(self, pipeline_id: int) -> dict[str, Any] | None: - """ - Get only the JSON config for a specific pipeline (for display/clone). - - Args: - pipeline_id: ID of the pipeline - - Returns: - Dictionary with pipeline_config_json, or None if not found - """ - import json - - query = """ - SELECT pipeline_config_json, context_id - FROM pipelines - WHERE pipeline_id = ? - """ - result = self.db.execute(query, (pipeline_id,), fetch=True) - if result and result[0][0]: - json_config = json.loads(result[0][0]) - return { - "pipeline_config_json": json_config, - "context_id": result[0][1] - } - return None - - def get_all_pipelines(self) -> list[dict[str, Any]]: - """ - Get all pipelines, ordered by pipeline ID. - - Returns: - List of pipeline dictionaries (ordered by pipeline_id DESC) - """ - query = """ - SELECT pipeline_id, context_id, pipeline_type, - pipeline_config, pipeline_config_json, flattened_config, status, current_shard_id, - shards_completed, total_samples_processed, metric_run_id, error, created_at - FROM pipelines - ORDER BY pipeline_id DESC - """ - result = self.db.execute(query, fetch=True) - pipelines = [] - if result: - import json - - for row in result: - # Decode the pipeline config from the database (use pipeline_config column) - decoded_config = decode_db_payload(row[3]) if row[3] else None - # Parse JSON config for display/analytics - json_config = json.loads(row[4]) if row[4] else None - flattened_config = json.loads(row[5]) if row[5] else {} - pipelines.append( - { - "pipeline_id": row[0], - "context_id": row[1], - "pipeline_type": row[2], - "pipeline_config": decoded_config, # Use decoded config for actual pipeline object - "pipeline_config_json": json_config, # JSON version for display/analytics - "flattened_config": flattened_config, # Flattened configuration for IC Ops panel display - "status": row[6], - "current_shard_id": row[7], - "shards_completed": row[8], - "total_samples_processed": row[9], - "metric_run_id": row[10], - "error": row[11], - "created_at": row[12], - } - ) - return pipelines - - def set_pipeline_status(self, pipeline_id: int, status: PipelineStatus): - """ - Update pipeline status. - - Args: - pipeline_id: ID of the pipeline - status: New status (PipelineStatus enum) - """ - query = "UPDATE pipelines SET status = ? WHERE pipeline_id = ?" - self.db.execute(query, (status.value, pipeline_id), commit=True) - - def set_pipeline_progress( - self, - pipeline_id: int, - current_shard_id: int, - shards_completed: int, - total_samples_processed: int, - ): - """ - Update pipeline progress metrics. - - Args: - pipeline_id: ID of the pipeline - current_shard_id: Current shard being processed - shards_completed: Number of shards completed - total_samples_processed: Total number of samples processed - """ - query = """ - UPDATE pipelines - SET current_shard_id = ?, shards_completed = ?, total_samples_processed = ? - WHERE pipeline_id = ? - """ - self.db.execute( - query, - (current_shard_id, shards_completed, total_samples_processed, pipeline_id), - commit=True, - ) - - def set_pipeline_current_shard(self, pipeline_id: int, shard_id: int): - """ - Update the current shard being processed by a pipeline. - - Args: - pipeline_id: ID of the pipeline - shard_id: Current shard ID being processed - """ - query = "UPDATE pipelines SET current_shard_id = ? WHERE pipeline_id = ?" - self.db.execute(query, (shard_id, pipeline_id), commit=True) - - def set_pipeline_error(self, pipeline_id: int, error: str): - """ - Set error message for a pipeline. - - Args: - pipeline_id: ID of the pipeline - error: Error message - """ - query = "UPDATE pipelines SET error = ? WHERE pipeline_id = ?" - self.db.execute(query, (error, pipeline_id), commit=True) - - def set_pipeline_metric_run_id(self, pipeline_id: int, metric_run_id: str): - """ - Set MetricLogger run ID for a pipeline. - - Args: - pipeline_id: ID of the pipeline - metric_run_id: MetricLogger run ID - """ - query = "UPDATE pipelines SET metric_run_id = ? WHERE pipeline_id = ?" - self.db.execute(query, (metric_run_id, pipeline_id), commit=True) - - # ============================================================================ - # ACTOR_TASKS TABLE METHODS - # ============================================================================ - - def create_actor_task( - self, - pipeline_id: int, - actor_id: int, - shard_id: int, - status: TaskStatus = TaskStatus.SCHEDULED, - ) -> int: - """ - Create a new actor task record. - - Args: - pipeline_id: ID of the pipeline - actor_id: ID of the actor - shard_id: ID of the shard - status: Initial status (default: TaskStatus.SCHEDULED) - - Returns: - task_id of the created task - """ - query = """ - INSERT INTO actor_tasks ( - pipeline_id, actor_id, shard_id, status, error_message - ) VALUES (?, ?, ?, ?, '') - """ - self.db.execute( - query, (pipeline_id, actor_id, shard_id, status.value), commit=True - ) - return self.db.cursor.lastrowid - - def get_actor_task(self, task_id: int) -> dict[str, Any] | None: - """ - Get actor task by ID. - - Args: - task_id: ID of the task - - Returns: - Dictionary with all task fields, or None if not found - """ - query = """ - SELECT task_id, pipeline_id, actor_id, shard_id, status, - error_message, started_at, completed_at, duration_seconds - FROM actor_tasks - WHERE task_id = ? - """ - result = self.db.execute(query, (task_id,), fetch=True) - if result: - row = result[0] - return { - "task_id": row[0], - "pipeline_id": row[1], - "actor_id": row[2], - "shard_id": row[3], - "status": row[4], - "error_message": row[5], - "started_at": row[6], - "completed_at": row[7], - "duration_seconds": row[8], - } - return None - - def get_running_actor_tasks(self) -> list[dict[str, Any]]: - """ - Get all currently running actor tasks. - - Returns: - List of running task dictionaries (ordered by task_id DESC) - """ - query = """ - SELECT task_id, pipeline_id, actor_id, shard_id, status, - error_message, started_at, completed_at, duration_seconds - FROM actor_tasks - WHERE status = ? - ORDER BY task_id DESC - """ - result = self.db.execute(query, (TaskStatus.IN_PROGRESS.value,), fetch=True) - tasks = [] - if result: - for row in result: - tasks.append( - { - "task_id": row[0], - "pipeline_id": row[1], - "actor_id": row[2], - "shard_id": row[3], - "status": row[4], - "error_message": row[5], - "started_at": row[6], - "completed_at": row[7], - "duration_seconds": row[8], - } - ) - return tasks - - def set_actor_task_status(self, task_id: int, status: TaskStatus): - """ - Update actor task status. - - Args: - task_id: ID of the task - status: New status (TaskStatus enum) - """ - query = "UPDATE actor_tasks SET status = ? WHERE task_id = ?" - self.db.execute(query, (status.value, task_id), commit=True) - - def set_actor_task_start_time(self, task_id: int, start_time: str): - """ - Set start time for an actor task. - - Args: - task_id: ID of the task - start_time: Start timestamp (use datetime.now().isoformat()) - """ - query = "UPDATE actor_tasks SET started_at = ? WHERE task_id = ?" - self.db.execute(query, (start_time, task_id), commit=True) - - def set_actor_task_end_time( - self, task_id: int, end_time: str, duration_seconds: float - ): - """ - Set end time and duration for an actor task. - - Args: - task_id: ID of the task - end_time: End timestamp (use datetime.now().isoformat()) - duration_seconds: Duration in seconds - """ - query = "UPDATE actor_tasks SET completed_at = ?, duration_seconds = ? WHERE task_id = ?" - self.db.execute(query, (end_time, duration_seconds, task_id), commit=True) - - def set_actor_task_error(self, task_id: int, error_message: str): - """ - Set error message for an actor task. - - Args: - task_id: ID of the task - error_message: Error message - """ - query = "UPDATE actor_tasks SET error_message = ? WHERE task_id = ?" - self.db.execute(query, (error_message, task_id), commit=True) - - # ============================================================================ - # INTERACTIVE_CONTROL TABLE METHODS - # ============================================================================ - - def create_ic_operation( - self, - operation: str, - pipeline_id: int = None, - request_data: str = None, - ) -> int: - """ - Create a new interactive control operation. - - Args: - operation: Type of operation ('stop', 'resume', 'delete', 'clone') - pipeline_id: ID of the pipeline (None for clone operation) - request_data: JSON string with operation data (e.g., model_config for clone) - - Returns: - ic_id of the created operation - """ - import time - - query = """ - INSERT INTO interactive_control ( - pipeline_id, operation, status, request_data, error, created_at - ) VALUES (?, ?, ?, ?, '', ?) - """ - from rapidfireai.evals.utils.constants import ICStatus - - self.db.execute( - query, - (pipeline_id, operation, ICStatus.PENDING.value, request_data, time.time()), - commit=True, - ) - return self.db.cursor.lastrowid - - def get_pending_ic_operations(self) -> list[dict[str, Any]]: - """ - Get all pending IC operations. - - Returns: - List of dictionaries with IC operation fields - """ - from rapidfireai.evals.utils.constants import ICStatus - - query = """ - SELECT ic_id, pipeline_id, operation, status, request_data, error, created_at, processed_at - FROM interactive_control - WHERE status = ? - ORDER BY created_at ASC - """ - results = self.db.execute(query, (ICStatus.PENDING.value,), fetch=True) - operations = [] - for row in results: - operations.append( - { - "ic_id": row[0], - "pipeline_id": row[1], - "operation": row[2], - "status": row[3], - "request_data": row[4], - "error": row[5], - "created_at": row[6], - "processed_at": row[7], - } - ) - return operations - - def get_ic_operation(self, ic_id: int) -> dict[str, Any] | None: - """ - Get IC operation by ID. - - Args: - ic_id: ID of the IC operation - - Returns: - Dictionary with IC operation fields, or None if not found - """ - query = """ - SELECT ic_id, pipeline_id, operation, status, request_data, error, created_at, processed_at - FROM interactive_control - WHERE ic_id = ? - """ - result = self.db.execute(query, (ic_id,), fetch=True) - if result: - row = result[0] - return { - "ic_id": row[0], - "pipeline_id": row[1], - "operation": row[2], - "status": row[3], - "request_data": row[4], - "error": row[5], - "created_at": row[6], - "processed_at": row[7], - } - return None - - def update_ic_operation_status(self, ic_id: int, status: str, error: str = ""): - """ - Update IC operation status. - - Args: - ic_id: ID of the IC operation - status: New status ('pending', 'processing', 'completed', 'failed') - error: Error message (if status is 'failed') - """ - import time - - query = """ - UPDATE interactive_control - SET status = ?, error = ?, processed_at = ? - WHERE ic_id = ? - """ - self.db.execute(query, (status, error, time.time(), ic_id), commit=True) - - def get_all_ic_operations(self) -> list[dict[str, Any]]: - """ - Get all IC operations, ordered by creation time. - - Returns: - List of dictionaries with IC operation fields - """ - query = """ - SELECT ic_id, pipeline_id, operation, status, request_data, error, created_at, processed_at - FROM interactive_control - ORDER BY created_at DESC - """ - results = self.db.execute(query, fetch=True) - operations = [] - for row in results: - operations.append( - { - "ic_id": row[0], - "pipeline_id": row[1], - "operation": row[2], - "status": row[3], - "request_data": row[4], - "error": row[5], - "created_at": row[6], - "processed_at": row[7], - } - ) - return operations - - -# Export for external use -__all__ = ["RFDatabase"] \ No newline at end of file diff --git a/rapidfireai/evals/db/tables.sql b/rapidfireai/evals/db/tables.sql deleted file mode 100644 index 8c245566..00000000 --- a/rapidfireai/evals/db/tables.sql +++ /dev/null @@ -1,94 +0,0 @@ --- RF-Inferno Database Schema --- SQLite schema for tracking multi-pipeline inference experiments --- Simplified: JSON-only storage, in-memory scheduler, essential state tracking - --- ============================================================================ --- EXPERIMENTS TABLE --- Tracks high-level experiment configuration and status --- ============================================================================ -CREATE TABLE IF NOT EXISTS experiments ( - experiment_id INTEGER PRIMARY KEY AUTOINCREMENT, - experiment_name TEXT NOT NULL, - metric_experiment_id TEXT, - num_shards INTEGER DEFAULT 0, - num_actors INTEGER NOT NULL, - num_cpus INTEGER, - num_gpus INTEGER, - status TEXT NOT NULL, -- 'running', 'completed', 'failed' - error TEXT DEFAULT '', - created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP -); - --- ============================================================================ --- CONTEXTS TABLE --- Stores unique RAG/context generation configurations --- Multiple pipelines can share the same context via context_id --- ============================================================================ -CREATE TABLE IF NOT EXISTS contexts ( - context_id INTEGER PRIMARY KEY AUTOINCREMENT, - context_hash TEXT NOT NULL UNIQUE, -- SHA256 hash for deduplication - rag_config_json TEXT, -- Full RAG configuration as JSON - prompt_config_json TEXT, -- Full prompt manager configuration as JSON - status TEXT NOT NULL, -- 'new', 'ongoing', 'deleted', 'failed' - error TEXT DEFAULT '', - started_at TIMESTAMP, - completed_at TIMESTAMP, - duration_seconds REAL -); - --- ============================================================================ --- PIPELINES TABLE --- Stores configuration for each inference pipeline --- ============================================================================ -CREATE TABLE IF NOT EXISTS pipelines ( - pipeline_id INTEGER PRIMARY KEY AUTOINCREMENT, - context_id INTEGER, -- Foreign key to contexts table (NULL if no RAG) - pipeline_type TEXT NOT NULL, -- 'vllm', 'openai_api', etc. - pipeline_config TEXT NOT NULL, -- Full pipeline configuration (encoded with encode_payload - includes functions/classes) - pipeline_config_json TEXT, -- JSON-serializable pipeline configuration (for analytics/display, excludes functions/classes) - flattened_config TEXT DEFAULT '{}', -- Flattened configuration for IC Ops panel display - status TEXT NOT NULL, -- 'new', 'ongoing', 'completed', 'stopped', 'deleted', 'failed' - current_shard_id INTEGER DEFAULT 0, -- Next shard to process - shards_completed INTEGER DEFAULT 0, -- Number of shards completed - total_samples_processed INTEGER DEFAULT 0, - metric_run_id TEXT, -- MetricLogger run ID for this pipeline - error TEXT DEFAULT '', - created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - - FOREIGN KEY (context_id) REFERENCES contexts(context_id) ON DELETE SET NULL -); - --- ============================================================================ --- ACTOR_TASKS TABLE --- Tracks individual tasks assigned to query processing actors --- ============================================================================ -CREATE TABLE IF NOT EXISTS actor_tasks ( - task_id INTEGER PRIMARY KEY AUTOINCREMENT, - pipeline_id INTEGER NOT NULL, - actor_id INTEGER NOT NULL, - shard_id INTEGER NOT NULL, - status TEXT NOT NULL, -- 'scheduled', 'in_progress', 'completed', 'failed' - error_message TEXT DEFAULT '', - started_at TIMESTAMP, - completed_at TIMESTAMP, - duration_seconds REAL, - - FOREIGN KEY (pipeline_id) REFERENCES pipelines(pipeline_id) ON DELETE CASCADE -); - --- ============================================================================ --- INTERACTIVE_CONTROL TABLE --- Tracks user-initiated dynamic pipeline control operations --- ============================================================================ -CREATE TABLE IF NOT EXISTS interactive_control ( - ic_id INTEGER PRIMARY KEY AUTOINCREMENT, - pipeline_id INTEGER, -- NULL for CLONE operation (creates new pipeline) - operation TEXT NOT NULL, -- 'stop', 'resume', 'delete', 'clone' - status TEXT NOT NULL, -- 'pending', 'processing', 'completed', 'failed' - request_data TEXT, -- JSON data for operation (e.g., model_config for clone) - error TEXT DEFAULT '', - created_at REAL NOT NULL, -- Unix timestamp - processed_at REAL, -- Unix timestamp when operation was processed - - FOREIGN KEY (pipeline_id) REFERENCES pipelines(pipeline_id) ON DELETE CASCADE -); diff --git a/rapidfireai/evals/dispatcher/README.md b/rapidfireai/evals/dispatcher/README.md deleted file mode 100644 index 053f8346..00000000 --- a/rapidfireai/evals/dispatcher/README.md +++ /dev/null @@ -1,239 +0,0 @@ -# RF-Inferno Dispatcher - -REST API for interactive control of running experiments. Allows dynamic pipeline management during execution. - -## Features - -- **Stop Pipeline**: Pause execution, can be resumed later -- **Resume Pipeline**: Continue stopped pipeline from last completed shard -- **Delete Pipeline**: Permanently remove pipeline -- **Clone Pipeline**: Create new pipeline using existing context - -## Automatic Startup - -The dispatcher **automatically starts** when you create an `Experiment` object. It runs in a background thread and cleans up automatically when the experiment ends. - -```python -from rapidfireai import Experiment - -# Dispatcher starts automatically on http://127.0.0.1:8851 (local) -# or via Colab proxy URL (in Google Colab) -experiment = Experiment(experiment_name="my-experiment", mode="evals") - -# Now you can use the dispatcher API for interactive control -# ... - -# Dispatcher automatically stops when experiment ends -experiment.end() -``` - -## Environment Support - -The dispatcher works in multiple environments: -- **Local Python** - Uses `http://127.0.0.1:8851` -- **Jupyter Notebook** - Uses `http://127.0.0.1:8851` -- **Google Colab** - Automatically uses Colab proxy URL with authentication - -### Google Colab - -When running in Google Colab, the dispatcher automatically: -- Detects the Colab environment (not triggered in regular Jupyter notebooks) -- Uses Colab's proxy URL instead of localhost -- Adds required authorization headers for proxy authentication - -The `NotebookUI` handles this automatically. For manual API requests, use helper functions: - -```python -import requests -from rapidfireai.evals.utils.constants import get_dispatcher_url, get_dispatcher_headers - -# Automatically returns correct URL (Colab proxy or localhost) and headers (with auth if needed) -dispatcher_url = get_dispatcher_url() -headers = get_dispatcher_headers() - -# Works in Colab, Jupyter, and local environments -response = requests.get( - f"{dispatcher_url}/dispatcher/health-check", - headers=headers -) -print(response.json()) -``` - -**Note:** Authorization headers are only included when running in Google Colab, not in regular Jupyter notebooks. - -## Manual Standalone Mode - -For testing or standalone use: -```bash -python -m rapidfireai.evals.dispatcher.dispatcher -``` -Runs on `http://127.0.0.1:8851` - -## API Endpoints - -Base URL: `http://127.0.0.1:8851/dispatcher` - -### Health Check -```bash -GET /dispatcher/health-check -``` - -### Stop Pipeline -```bash -POST /dispatcher/stop-pipeline -Content-Type: application/json - -{ - "pipeline_id": 123 -} -``` - -**Response:** -```json -{ - "ic_id": 1, - "message": "Stop request created for pipeline 123" -} -``` - -### Resume Pipeline -```bash -POST /dispatcher/resume-pipeline -Content-Type: application/json - -{ - "pipeline_id": 123 -} -``` - -### Delete Pipeline -```bash -POST /dispatcher/delete-pipeline -Content-Type: application/json - -{ - "pipeline_id": 123 -} -``` - -### Clone Pipeline (VLLM) -```bash -POST /dispatcher/clone-pipeline -Content-Type: application/json - -{ - "context_id": 1, - "pipeline_name": "my_cloned_vllm_pipeline", - "pipeline_type": "vllm", - "model_config": { - "model": "Qwen/Qwen2.5-3B-Instruct", - "dtype": "half", - "gpu_memory_utilization": 0.6 - }, - "sampling_params": { - "temperature": 0.8, - "top_p": 0.95, - "max_tokens": 512 - } -} -``` - -### Clone Pipeline (OpenAI) -```bash -POST /dispatcher/clone-pipeline -Content-Type: application/json - -{ - "context_id": 1, - "pipeline_name": "my_cloned_openai_pipeline", - "pipeline_type": "openai", - "client_config": { - "api_key": "sk-...", - "base_url": "https://api.openai.com/v1" - }, - "model_config": { - "model": "gpt-4", - "temperature": 0.7, - "top_p": 0.9 - }, - "rpm_limit": 500, - "tpm_limit": 90000 -} -``` - -### Check Operation Status -```bash -GET /dispatcher/operation-status/{ic_id} -``` - -**Response:** -```json -{ - "ic_id": 1, - "pipeline_id": 123, - "operation": "stop", - "status": "completed", - "error": "", - "created_at": 1698765432.123, - "processed_at": 1698765435.456 -} -``` - -### Get All Operations -```bash -GET /dispatcher/all-operations -``` - -## Usage Example (Python) - -```python -import requests - -# Stop a pipeline -response = requests.post( - "http://127.0.0.1:8851/dispatcher/stop-pipeline", - json={"pipeline_id": 123} -) -ic_id = response.json()["ic_id"] - -# Check status -status = requests.get( - f"http://127.0.0.1:8851/dispatcher/operation-status/{ic_id}" -) -print(status.json()) - -# Clone a pipeline -response = requests.post( - "http://127.0.0.1:8851/dispatcher/clone-pipeline", - json={ - "context_id": 1, - "pipeline_name": "cloned_pipeline", - "model_config": {...}, - "sampling_params": {...} - } -) -``` - -## Operation Lifecycle - -1. **User sends request** โ†’ Dispatcher validates and creates IC operation (status: `pending`) -2. **Controller polls database** โ†’ Finds pending operation (status: `processing`) -3. **Controller executes operation** โ†’ Modifies scheduler, updates database -4. **Operation completes** โ†’ Status changes to `completed` or `failed` - -## Error Handling - -All endpoints return JSON with error details: -```json -{ - "error": "Pipeline 123 not found", - "traceback": "..." -} -``` - -HTTP Status Codes: -- `200`: Success -- `400`: Bad request (missing parameters) -- `404`: Resource not found -- `500`: Server error - diff --git a/rapidfireai/evals/dispatcher/__init__.py b/rapidfireai/evals/dispatcher/__init__.py deleted file mode 100644 index 8939adaf..00000000 --- a/rapidfireai/evals/dispatcher/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -""" -Dispatcher module for interactive control of RF-Inferno experiments. - -Provides REST API for dynamic pipeline management: -- Stop, Resume, Delete pipelines -- Clone new pipelines from existing contexts - -The dispatcher automatically starts when an Experiment is created and runs -in a background thread. It can also be run standalone for testing. -""" - -from rapidfireai.evals.dispatcher.dispatcher import Dispatcher, run_dispatcher, start_dispatcher_thread - -__all__ = ["Dispatcher", "run_dispatcher", "start_dispatcher_thread"] diff --git a/rapidfireai/evals/dispatcher/dispatcher.py b/rapidfireai/evals/dispatcher/dispatcher.py deleted file mode 100644 index 0e0ec1aa..00000000 --- a/rapidfireai/evals/dispatcher/dispatcher.py +++ /dev/null @@ -1,1152 +0,0 @@ -""" -Dispatcher REST API for Interactive Control of RF-Inferno Experiments. - -Provides HTTP endpoints for dynamic pipeline management during experiment execution. -FIXED: Now properly handles CORS preflight (OPTIONS) requests for VS Code/Cursor webview. -""" - -import json -import logging -import os -import threading -import traceback -from flask import Flask, Response, jsonify, request -from flask_cors import CORS -from waitress import serve - -from rapidfireai.evals.db import RFDatabase -from rapidfireai.evals.utils.logger import RFLogger, SafeLoggerAdapter -from rapidfireai.utils.constants import DispatcherConfig, ColabConfig, RF_LOG_PATH, RF_LOG_FILENAME -from rapidfireai.utils.dispatcher_utils import check_experiment_running -from rapidfireai.evals.utils.constants import ICOperation - -CORS_ALLOWED_ORIGINS = "*" # Allow all origins - -class Dispatcher: - """ - REST API server for interactive control of running experiments. - - Handles user requests to: - - Stop pipelines (pause execution, can be resumed) - - Resume pipelines (continue from where stopped) - - Delete pipelines (permanently remove) - - Clone pipelines (create new pipeline with existing context) - - Also provides frontend-compatible "run" aliases for pipeline endpoints - to support the MLflow dashboard IC Ops panel. - """ - - # Status mapping from evals (lowercase) to fit (capitalized) format - STATUS_MAP = { - "new": "New", - "ongoing": "Ongoing", - "completed": "Completed", - "stopped": "Stopped", - "deleted": "Deleted", - "failed": "Failed", - } - - def __init__(self) -> None: - """Initialize the Dispatcher with database connection and Flask app.""" - # initialize loggers - self.logger_experiment_name: str | None = None - self.logger: SafeLoggerAdapter | None = None - - # Create database handle - self.db: RFDatabase = RFDatabase() - - # Create Flask app - self.app: Flask = Flask(__name__) - - # Enable CORS for local development - # Dispatcher runs on localhost, safe to allow all origins - # supports_credentials=True is required for Colab proxy auth (credentials: 'include' in JS) - _ = CORS( - self.app, - resources={ - r"/*": { - "origins": CORS_ALLOWED_ORIGINS, - "methods": ["GET", "POST", "PUT", "DELETE", "OPTIONS"], - "allow_headers": ["Content-Type", "Authorization"], - "expose_headers": ["Content-Type"], - "supports_credentials": True if ColabConfig.ON_COLAB else False, - } - }, - ) - - # Register routes - self.register_routes() - - def _get_logger(self) -> SafeLoggerAdapter: - """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(experiment_name=current_experiment_name).get_logger("dispatcher") - self.logger_experiment_name = current_experiment_name - return self.logger - - def _transform_pipeline_to_run(self, pipeline: dict) -> dict: - """ - Transform pipeline data to run format for frontend compatibility. - - The MLflow frontend expects run data in a specific format that differs - from the evals pipeline format. This method converts pipeline data - to match the expected run format. - - Args: - pipeline: Pipeline data from database - - Returns: - Run-formatted data for frontend consumption - """ - status = pipeline.get("status", "") - return { - "run_id": pipeline.get("pipeline_id"), - "status": self.STATUS_MAP.get(status.lower(), status), - "metric_run_id": pipeline.get("metric_run_id"), - "config": pipeline.get("pipeline_config_json", {}), - "flattened_config": pipeline.get("flattened_config", {}), - "num_chunks_visited": pipeline.get("shards_completed", 0), - "total_samples_processed": pipeline.get("total_samples_processed", 0), - "error": pipeline.get("error", ""), - } - - def register_routes(self) -> None: - """Register all REST API routes with OPTIONS support for CORS preflight.""" - route_prefix = "/dispatcher" - - # CRITICAL: Add before_request handler to handle OPTIONS preflight requests globally - @self.app.before_request - def handle_preflight(): - if request.method == "OPTIONS": - response = jsonify({}) - response.headers.add("Access-Control-Allow-Origin", CORS_ALLOWED_ORIGINS) - response.headers.add("Access-Control-Allow-Headers", "Content-Type,Authorization") - response.headers.add("Access-Control-Allow-Methods", "GET,PUT,POST,DELETE,OPTIONS") - response.headers.add("Access-Control-Max-Age", "3600") - return response - - # Health check - self.app.add_url_rule( - f"{route_prefix}/health-check", "health_check", self.health_check, methods=["GET", "OPTIONS"] - ) - - # Interactive control operations - self.app.add_url_rule( - f"{route_prefix}/stop-pipeline", "stop_pipeline", self.stop_pipeline, methods=["POST", "OPTIONS"] - ) - self.app.add_url_rule( - f"{route_prefix}/resume-pipeline", "resume_pipeline", self.resume_pipeline, methods=["POST", "OPTIONS"] - ) - self.app.add_url_rule( - f"{route_prefix}/delete-pipeline", "delete_pipeline", self.delete_pipeline, methods=["POST", "OPTIONS"] - ) - self.app.add_url_rule( - f"{route_prefix}/clone-pipeline", "clone_pipeline", self.clone_pipeline, methods=["POST", "OPTIONS"] - ) - - # Status queries - self.app.add_url_rule( - f"{route_prefix}/operation-status/", - "get_operation_status", - self.get_operation_status, - methods=["GET", "OPTIONS"], - ) - self.app.add_url_rule( - f"{route_prefix}/all-operations", "get_all_operations", self.get_all_operations, methods=["GET", "OPTIONS"] - ) - - # Pipeline queries (for UI) - self.app.add_url_rule( - f"{route_prefix}/list-all-pipeline-ids", - "list_all_pipeline_ids", - self.list_all_pipeline_ids, - methods=["GET", "OPTIONS"], - ) - self.app.add_url_rule( - f"{route_prefix}/get-pipeline-config-json/", - "get_pipeline_config_json", - self.get_pipeline_config_json, - methods=["GET", "OPTIONS"], - ) - # Legacy endpoints (kept for backwards compatibility) - self.app.add_url_rule( - f"{route_prefix}/get-all-pipelines", "get_all_pipelines", self.get_all_pipelines, methods=["GET", "OPTIONS"] - ) - self.app.add_url_rule( - f"{route_prefix}/get-pipeline", "get_pipeline", self.get_pipeline, methods=["POST", "OPTIONS"] - ) - - # ============================================================ - # Frontend-compatible "run" aliases (for MLflow dashboard IC Ops) - # These endpoints mirror the fit dispatcher API to allow the - # frontend to work with evals mode without modifications. - # ============================================================ - - # Run query aliases - self.app.add_url_rule( - f"{route_prefix}/get-all-runs", "get_all_runs", self.get_all_runs, methods=["GET", "OPTIONS"] - ) - self.app.add_url_rule( - f"{route_prefix}/get-run", "get_run", self.get_run, methods=["POST", "OPTIONS"] - ) - - # IC Ops aliases (run โ†’ pipeline) - self.app.add_url_rule( - f"{route_prefix}/stop-run", "stop_run", self.stop_run, methods=["POST", "OPTIONS"] - ) - self.app.add_url_rule( - f"{route_prefix}/resume-run", "resume_run", self.resume_run, methods=["POST", "OPTIONS"] - ) - self.app.add_url_rule( - f"{route_prefix}/delete-run", "delete_run", self.delete_run, methods=["POST", "OPTIONS"] - ) - self.app.add_url_rule( - f"{route_prefix}/clone-modify-run", "clone_modify_run", self.clone_modify_run, methods=["POST", "OPTIONS"] - ) - - # Experiment info endpoints (for frontend context) - self.app.add_url_rule( - f"{route_prefix}/get-running-experiment", - "get_running_experiment", - self.get_running_experiment, - methods=["GET", "OPTIONS"], - ) - self.app.add_url_rule( - f"{route_prefix}/get-all-experiment-names", - "get_all_experiment_names", - self.get_all_experiment_names, - methods=["GET", "OPTIONS"], - ) - self.app.add_url_rule( - f"{route_prefix}/is-experiment-running", - "is_experiment_running", - self.is_experiment_running, - methods=["POST", "OPTIONS"], - ) - - # Logging endpoints (placeholders for frontend compatibility) - self.app.add_url_rule( - f"{route_prefix}/get-experiment-logs", - "get_experiment_logs", - self.get_experiment_logs, - methods=["POST", "OPTIONS"], - ) - self.app.add_url_rule( - f"{route_prefix}/get-ic-logs", "get_ic_logs", self.get_ic_logs, methods=["POST", "OPTIONS"] - ) - - def health_check(self) -> tuple[Response, int]: - """Health check endpoint.""" - # Handle OPTIONS preflight - if request.method == "OPTIONS": - return jsonify({"status": "ok"}), 200 - - try: - return jsonify({"status": "ok", "message": "Dispatcher is running"}), 200 - except Exception as e: - return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 - - def stop_pipeline(self) -> tuple[Response, int]: - """ - Stop a running pipeline. - - Request body: - { - "pipeline_id": int - } - - Returns: - ic_id of the created operation - """ - # Handle OPTIONS preflight - if request.method == "OPTIONS": - return jsonify({}), 200 - - try: - data = request.get_json() - if not data: - return jsonify({"error": "No JSON data provided"}), 400 - - pipeline_id = data.get("pipeline_id") - if pipeline_id is None: - return jsonify({"error": "pipeline_id is required"}), 400 - - # Validate pipeline exists - pipeline = self.db.get_pipeline(pipeline_id) - if not pipeline: - return jsonify({"error": f"Pipeline {pipeline_id} not found"}), 404 - - # Create IC operation - ic_id = self.db.create_ic_operation( - operation=ICOperation.STOP.value, - pipeline_id=pipeline_id, - ) - - return jsonify({"ic_id": ic_id, "message": f"Stop request created for pipeline {pipeline_id}"}), 200 - - except Exception as e: - return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 - - def resume_pipeline(self) -> tuple[Response, int]: - """ - Resume a stopped pipeline. - - Request body: - { - "pipeline_id": int - } - - Returns: - ic_id of the created operation - """ - # Handle OPTIONS preflight - if request.method == "OPTIONS": - return jsonify({}), 200 - - try: - data = request.get_json() - if not data: - return jsonify({"error": "No JSON data provided"}), 400 - - pipeline_id = data.get("pipeline_id") - if pipeline_id is None: - return jsonify({"error": "pipeline_id is required"}), 400 - - # Validate pipeline exists - pipeline = self.db.get_pipeline(pipeline_id) - if not pipeline: - return jsonify({"error": f"Pipeline {pipeline_id} not found"}), 404 - - # Create IC operation - ic_id = self.db.create_ic_operation( - operation=ICOperation.RESUME.value, - pipeline_id=pipeline_id, - ) - - return jsonify({"ic_id": ic_id, "message": f"Resume request created for pipeline {pipeline_id}"}), 200 - - except Exception as e: - return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 - - def delete_pipeline(self) -> tuple[Response, int]: - """ - Delete a pipeline permanently. - - Request body: - { - "pipeline_id": int - } - - Returns: - ic_id of the created operation - """ - # Handle OPTIONS preflight - if request.method == "OPTIONS": - return jsonify({}), 200 - - try: - data = request.get_json() - if not data: - return jsonify({"error": "No JSON data provided"}), 400 - - pipeline_id = data.get("pipeline_id") - if pipeline_id is None: - return jsonify({"error": "pipeline_id is required"}), 400 - - # Validate pipeline exists - pipeline = self.db.get_pipeline(pipeline_id) - if not pipeline: - return jsonify({"error": f"Pipeline {pipeline_id} not found"}), 404 - - # Create IC operation - ic_id = self.db.create_ic_operation( - operation=ICOperation.DELETE.value, - pipeline_id=pipeline_id, - ) - - return jsonify({"ic_id": ic_id, "message": f"Delete request created for pipeline {pipeline_id}"}), 200 - - except Exception as e: - return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 - - def clone_pipeline(self) -> tuple[Response, int]: - """ - Clone a new pipeline from a parent pipeline with edited configuration. - - The clone inherits the parent's context_id, RAG, and prompt_manager. - Only the JSON-editable parameters can be modified. - - Request body: - { - "parent_pipeline_id": int, # ID of the pipeline to clone - "config_json": { # Edited configuration - "pipeline_type": "vllm" | "openai", - "model_config": {...}, - "sampling_params": {...}, # for vLLM - "client_config": {...}, # for OpenAI - "batch_size": int, # optional - "online_strategy_kwargs": {...} # optional - } - } - - Returns: - ic_id of the created operation - """ - # Handle OPTIONS preflight - if request.method == "OPTIONS": - return jsonify({}), 200 - - try: - data = request.get_json() - if not data: - return jsonify({"error": "No JSON data provided"}), 400 - - parent_pipeline_id = data.get("parent_pipeline_id") - if parent_pipeline_id is None: - return jsonify({"error": "parent_pipeline_id is required"}), 400 - - config_json = data.get("config_json") - if not config_json: - return jsonify({"error": "config_json is required"}), 400 - - # Validate parent pipeline exists - parent_pipeline = self.db.get_pipeline(parent_pipeline_id) - if not parent_pipeline: - return jsonify({"error": f"Parent pipeline {parent_pipeline_id} not found"}), 404 - - # Validate config_json has required fields - pipeline_type = config_json.get("pipeline_type") - if not pipeline_type: - return jsonify({"error": "config_json must include 'pipeline_type'"}), 400 - - if pipeline_type.lower() not in ["vllm", "openai"]: - return jsonify({"error": "pipeline_type must be 'vllm' or 'openai'"}), 400 - - # Type-specific validation - if pipeline_type.lower() == "vllm": - if "model_config" not in config_json or "sampling_params" not in config_json: - return jsonify({"error": "vLLM pipelines require 'model_config' and 'sampling_params'"}), 400 - - elif pipeline_type.lower() == "openai": - if "client_config" not in config_json or "model_config" not in config_json: - return jsonify({"error": "OpenAI pipelines require 'client_config' and 'model_config'"}), 400 - - # Prepare request data for IC operation - request_data = { - "parent_pipeline_id": parent_pipeline_id, - "config_json": config_json, - } - - # Create IC operation (pipeline_id is None for CLONE, as new ID will be generated) - ic_id = self.db.create_ic_operation( - operation=ICOperation.CLONE.value, - pipeline_id=None, - request_data=json.dumps(request_data), - ) - - return ( - jsonify( - { - "ic_id": ic_id, - "message": f"Clone request created from parent pipeline {parent_pipeline_id}", - } - ), - 200, - ) - - except Exception as e: - return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 - - def get_operation_status(self, ic_id: int) -> tuple[Response, int]: - """ - Get status of a specific IC operation. - - Args: - ic_id: ID of the IC operation - - Returns: - Operation details including status - """ - # Handle OPTIONS preflight - if request.method == "OPTIONS": - return jsonify({}), 200 - - try: - operation = self.db.get_ic_operation(ic_id) - if not operation: - return jsonify({"error": f"Operation {ic_id} not found"}), 404 - - return jsonify(operation), 200 - - except Exception as e: - return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 - - def get_all_operations(self) -> tuple[Response, int]: - """ - Get all IC operations (for monitoring/debugging). - - Returns: - List of all operations - """ - # Handle OPTIONS preflight - if request.method == "OPTIONS": - return jsonify({}), 200 - - try: - operations = self.db.get_all_ic_operations() - return jsonify(operations), 200 - - except Exception as e: - return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 - - def list_all_pipeline_ids(self) -> tuple[Response, int]: - """ - Get lightweight list of pipeline IDs with minimal info (optimized for polling). - - Returns: - List of pipelines with only: pipeline_id, status, shards_completed, total_samples_processed - """ - # Handle OPTIONS preflight - if request.method == "OPTIONS": - return jsonify({}), 200 - - try: - pipelines = self.db.get_all_pipeline_ids() - return jsonify(pipelines), 200 - - except Exception as e: - return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 - - def get_pipeline_config_json(self, pipeline_id: int) -> tuple[Response, int]: - """ - Get only the config JSON for a specific pipeline (for clone operations). - - Args: - pipeline_id: ID of the pipeline (from URL path) - - Returns: - Pipeline config JSON - """ - # Handle OPTIONS preflight - if request.method == "OPTIONS": - return jsonify({}), 200 - - try: - config_data = self.db.get_pipeline_config_json(pipeline_id) - if not config_data: - return jsonify({"error": f"Pipeline {pipeline_id} not found"}), 404 - - return jsonify(config_data), 200 - - except Exception as e: - return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 - - def get_all_pipelines(self) -> tuple[Response, int]: - """ - Get all pipelines (for UI dropdown). - - LEGACY: Use list_all_pipeline_ids() for better performance. - - Returns: - List of all pipelines - """ - # Handle OPTIONS preflight - if request.method == "OPTIONS": - return jsonify({}), 200 - - try: - pipelines = self.db.get_all_pipelines() - return jsonify(pipelines), 200 - - except Exception as e: - return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 - - def get_pipeline(self) -> tuple[Response, int]: - """ - Get details of a specific pipeline. - - Request body: - { - "pipeline_id": int - } - - Returns: - Pipeline details - """ - # Handle OPTIONS preflight - if request.method == "OPTIONS": - return jsonify({}), 200 - - try: - data = request.get_json() - if not data: - return jsonify({"error": "No JSON data provided"}), 400 - - pipeline_id = data.get("pipeline_id") - if pipeline_id is None: - return jsonify({"error": "pipeline_id is required"}), 400 - - pipeline = self.db.get_pipeline(pipeline_id) - if not pipeline: - return jsonify({"error": f"Pipeline {pipeline_id} not found"}), 404 - - return jsonify(pipeline), 200 - - except Exception as e: - return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 - - # ============================================================ - # Frontend-compatible "run" alias handlers - # ============================================================ - - def get_all_runs(self) -> tuple[Response, int]: - """ - Get all runs (alias for get_all_pipelines with response transformation). - - This endpoint provides frontend compatibility by returning pipeline data - in the "run" format expected by the MLflow dashboard. - - Returns: - List of runs in frontend-compatible format - """ - if request.method == "OPTIONS": - return jsonify({}), 200 - - try: - pipelines = self.db.get_all_pipelines() - runs = [self._transform_pipeline_to_run(p) for p in pipelines] - return jsonify(runs), 200 - - except Exception as e: - return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 - - def get_run(self) -> tuple[Response, int]: - """ - Get a specific run (alias for get_pipeline with response transformation). - - Request body: - { - "run_id": int or str (pipeline_id or metricrun_id) - } - - Returns: - Run details in frontend-compatible format - """ - if request.method == "OPTIONS": - return jsonify({}), 200 - - try: - data = request.get_json() - if not data: - return jsonify({"error": "No JSON data provided"}), 400 - - run_id = data.get("run_id") - if run_id is None: - return jsonify({"error": "run_id is required"}), 400 - - - # Validate pipeline exists - pipeline = self.db.get_pipeline(run_id) - if not pipeline: - return jsonify({"error": f"Run {run_id} not found"}), 400 - - return jsonify(self._transform_pipeline_to_run(pipeline)), 200 - - except Exception as e: - return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 - - def stop_run(self) -> tuple[Response, int]: - """Stop the run""" - try: - data = request.get_json() - run_id = data["run_id"] - task = ICOperation.STOP - - # get pipeline from db - pipeline = self.db.get_pipeline(run_id) - - # create ic ops task - _ = self.db.create_ic_operation( - operation=task.value, - pipeline_id=pipeline["pipeline_id"], - ) - - # log the task - logger = self._get_logger() - logger.info(f"Received stop task for run_id: {run_id}") - return jsonify({}), 200 - except Exception as e: - logger = self._get_logger() - logger.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 = ICOperation.RESUME - - # get pipeline from db - pipeline = self.db.get_pipeline(run_id) - - # set resume run task - _ = self.db.create_ic_operation( - operation=task.value, - pipeline_id=pipeline["pipeline_id"], - ) - - # log the task - logger = self._get_logger() - logger.info(f"Received resume task for run_id: {run_id}") - return jsonify({}), 200 - except Exception as e: - logger = self._get_logger() - logger.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 = ICOperation.DELETE - - # get pipeline from db - pipeline = self.db.get_pipeline(run_id) - - # set delete run task - _ = self.db.create_ic_operation( - operation=task.value, - pipeline_id=pipeline["pipeline_id"], - ) - - # log the task - logger = self._get_logger() - logger.info(f"Received delete task for run_id: {run_id}") - return jsonify({}), 200 - except Exception as e: - logger = self._get_logger() - logger.error(f"Error in delete_run: {e}", exc_info=True) - return jsonify({"error": str(e) + " " + str(traceback.format_exc())}), 500 - - def clone_modify_run(self) -> tuple[Response, int]: - """ - Clone and modify a run (alias for clone_pipeline with payload transformation). - - Request body: - { - "run_id": int, # ID of the run to clone (maps to parent_pipeline_id) - "config": dict, # New configuration - "warm_start": bool # Optional, passed through but may not apply to evals - } - - Returns: - IC operation result - """ - if request.method == "OPTIONS": - return jsonify({}), 200 - - try: - data = request.get_json() - if not data: - return jsonify({"error": "No JSON data provided"}), 400 - - run_id = data.get("run_id") - if run_id is None: - return jsonify({"error": "run_id is required"}), 400 - - config = data.get("config", {}) - warm_start = data.get("warm_start", False) - - # Validate parent pipeline exists - parent_pipeline = self.db.get_pipeline(run_id) - if not parent_pipeline: - return jsonify({"error": f"Run {run_id} not found"}), 404 - - # Get the parent's config to use as base if config is empty - if not config: - config = parent_pipeline.get("pipeline_config_json", {}) - - # Prepare request data for IC operation - request_data = { - "parent_pipeline_id": run_id, - "config_json": config, - "warm_start": warm_start, # Passed through for potential future use - } - - # Create IC operation - ic_id = self.db.create_ic_operation( - operation=ICOperation.CLONE.value, - pipeline_id=None, - request_data=json.dumps(request_data), - ) - - return jsonify({"ic_id": ic_id, "message": f"Clone request created from run {run_id}"}), 200 - - except Exception as e: - return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 - - # ============================================================ - # Experiment info endpoints - # ============================================================ - - def get_running_experiment(self) -> tuple[Response, int]: - """ - Get the currently running experiment. - - Returns: - Experiment details including ID, name, status, and MLflow experiment ID - """ - if request.method == "OPTIONS": - return jsonify({}), 200 - - try: - experiment = self.db.get_running_experiment() - if not experiment: - return jsonify({"error": "No running experiment found"}), 404 - - return jsonify({ - "experiment_id": experiment.get("experiment_id"), - "experiment_name": experiment.get("experiment_name"), - "status": experiment.get("status"), - "metric_experiment_id": experiment.get("metric_experiment_id"), - }), 200 - - except Exception as e: - return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 - - def get_all_experiment_names(self) -> tuple[Response, int]: - """ - Get all experiment names. - - Returns: - List of experiment names - """ - if request.method == "OPTIONS": - return jsonify({}), 200 - - try: - names = self.db.get_all_experiment_names() - return jsonify(names), 200 - - except Exception as e: - return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 - - def is_experiment_running(self) -> tuple[Response, int]: - """ - Check if a specific experiment is currently running. - - Request body: - { - "experiment_name": str - The experiment name to check - } - - Returns: - { - "is_running": bool - True if the experiment is currently running - } - """ - if request.method == "OPTIONS": - return jsonify({}), 200 - - try: - data = request.get_json() - if not data or "experiment_name" not in data: - return jsonify({"error": "experiment_name is required"}), 400 - - is_running = check_experiment_running(self.db, data["experiment_name"]) - return jsonify({"is_running": is_running}), 200 - - except Exception: - # If anything fails, assume experiment is not running (safer to disable button) - return jsonify({"is_running": False}), 200 - - # ============================================================ - # Logging endpoints - # ============================================================ - - def get_experiment_logs(self) -> tuple[Response, int]: - """ - Get experiment logs. - - Reads logs from the experiment log file and returns entries matching - the experiment name. - - Request body: - { - "experiment_name": str (optional) - if not provided, uses running experiment - } - - Returns: - List of log entries for the experiment - """ - if request.method == "OPTIONS": - return jsonify({}), 200 - - 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: - running_exp = self.db.get_running_experiment() - if running_exp: - experiment_name = running_exp.get("experiment_name") - - if not experiment_name: - return jsonify([]), 200 - - log_file_path = os.path.join(RF_LOG_PATH, experiment_name, RF_LOG_FILENAME) - - if not os.path.exists(log_file_path): - return jsonify([]), 200 - - experiment_logs = [] - with open(log_file_path, encoding="utf-8") as f: - for line in f: - # Filter for lines containing the experiment name - if f"[{experiment_name}:" in line or f"| {experiment_name} |" in line: - experiment_logs.append(line.strip()) - - return jsonify(experiment_logs), 200 - except Exception as e: - return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 - - def get_ic_logs(self) -> tuple[Response, int]: - """ - Get interactive control logs. - - Reads logs from the experiment log file and returns entries related - to interactive control operations. - - Request body: - { - "experiment_name": str (optional) - if not provided, uses running experiment - } - - Returns: - List of IC log entries for the experiment - """ - if request.method == "OPTIONS": - return jsonify({}), 200 - - 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: - running_exp = self.db.get_running_experiment() - if running_exp: - experiment_name = running_exp.get("experiment_name") - - if not experiment_name: - return jsonify([]), 200 - - log_file_path = os.path.join(RF_LOG_PATH, experiment_name, RF_LOG_FILENAME) - - if not os.path.exists(log_file_path): - return jsonify([]), 200 - - # Read and filter logs for interactive-control entries - ic_logs = [] - with open(log_file_path, encoding="utf-8") as f: - for line in f: - if f"[{experiment_name}:InteractiveControl]" in line: - ic_logs.append(line.strip()) - - return jsonify(ic_logs), 200 - except Exception as e: - return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 - - -def run_dispatcher(host: str = "0.0.0.0", port: int = 8851) -> None: - """ - Run the dispatcher server. - - This function is designed to be called in a separate thread from the main experiment. - - Args: - host: Host to bind to (default: 0.0.0.0) - port: Port to bind to (default: 8851) - """ - try: - dispatcher = Dispatcher() - - # Enable verbose logging for waitress - logging.getLogger("waitress").setLevel(logging.WARNING) - - # Use waitress to serve the Flask app - serve(dispatcher.app, host=host, port=port, threads=6) - except Exception as e: - # Catch all exceptions to prevent thread crashes - print(f"CRITICAL: Dispatcher crashed: {e}") - traceback.print_exc() - - -# Global dispatcher thread tracking to avoid killing our own process -_dispatcher_thread: threading.Thread | None = None -_dispatcher_lock = threading.Lock() - - -def _check_port_in_use(host: str, port: int) -> bool: - """Check if a port is already in use.""" - import socket - - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - return s.connect_ex((host, port)) == 0 - - -def _is_dispatcher_thread_alive() -> bool: - """Check if our dispatcher thread is still running.""" - global _dispatcher_thread - return _dispatcher_thread is not None and _dispatcher_thread.is_alive() - - -def _cleanup_old_dispatcher(port: int, logger=None) -> None: - """ - Kill any old dispatcher processes using the port. - - IMPORTANT: Only kills external processes, not threads in the current process. - If the port is in use by our own dispatcher thread, this is a no-op. - """ - import os - import subprocess - - # If our dispatcher thread is alive, don't try to kill anything - if _is_dispatcher_thread_alive(): - msg = "Dispatcher thread is already running in this process, skipping cleanup" - if logger: - logger.debug(msg) - return - - current_pid = str(os.getpid()) - - try: - # Find process using the port - result = subprocess.run(["lsof", "-ti", f":{port}"], capture_output=True, text=True, timeout=2) - if result.returncode == 0 and result.stdout.strip(): - pids = result.stdout.strip().split("\n") - for pid in pids: - # CRITICAL: Never kill our own process - if pid.strip() == current_pid: - msg = f"Port {port} is used by current process (PID {pid}), skipping kill" - if logger: - logger.warning(msg) - else: - print(f"WARNING: {msg}") - continue - try: - subprocess.run(["kill", "-9", pid], timeout=2) - msg = f"Killed old process (PID {pid}) on port {port}" - if logger: - logger.info(msg) - else: - print(msg) - except: - pass - except: - pass # lsof might not be available - - -def start_dispatcher_thread(host: str = "0.0.0.0", port: int = 8851, logger=None) -> threading.Thread | None: - """ - Start the dispatcher REST API server in a background daemon thread. - - The dispatcher enables interactive control (stop/resume/delete/clone pipelines) - during experiment execution. It runs as a daemon thread and automatically - cleans up when the experiment ends. - - If a dispatcher thread is already running in this process, returns the existing - thread instead of starting a new one. - - Args: - host: Host to bind to (default: 0.0.0.0, localhost only) - port: Port to bind to (default: 8851) - logger: Optional logger instance for logging (if None, uses print) - - Returns: - The dispatcher thread object, or None if startup failed - """ - global _dispatcher_thread - - with _dispatcher_lock: - # Check if our dispatcher thread is already running - if _is_dispatcher_thread_alive(): - msg = f"Dispatcher thread already running on port {port}, reusing existing thread" - if logger: - logger.info(msg) - else: - print(msg) - return _dispatcher_thread - - try: - # Check if port is in use (by an external process) - if _check_port_in_use(host, port): - msg = f"Port {port} is already in use. Attempting cleanup..." - if logger: - logger.warning(msg) - else: - print(f"WARNING: {msg}") - - # Try to clean up old process (will skip if it's our own process) - _cleanup_old_dispatcher(port, logger) - - # Wait a moment and check again - import time - - time.sleep(0.5) - if _check_port_in_use(host, port): - error_msg = f"Port {port} still in use after cleanup. Restart your kernel or system." - if logger: - logger.error(error_msg) - else: - print(f"ERROR: {error_msg}") - return None - - # Create daemon thread (auto-cleanup when main process ends) - _dispatcher_thread = threading.Thread( - target=run_dispatcher, kwargs={"host": host, "port": port}, daemon=True, name="DispatcherThread" - ) - _dispatcher_thread.start() - - msg = f"Started interactive control dispatcher on http://{host}:{port}" - if logger: - logger.info(msg) - logger.info("Use dispatcher API to dynamically stop/resume/delete/clone pipelines") - else: - print(msg) - print("Use dispatcher API to dynamically stop/resume/delete/clone pipelines") - - return _dispatcher_thread - - except Exception as e: - error_msg = f"Failed to start dispatcher: {e}. Interactive control will not be available." - if logger: - logger.warning(error_msg) - else: - print(f"WARNING: {error_msg}") - return None - - -def serve_forever() -> Flask: - """Start the Dispatcher via Gunicorn. - - This function is the entry point for Gunicorn WSGI server. - It returns the Flask app instance for Gunicorn to serve. - - Returns: - Flask app instance - """ - return Dispatcher().app - - -if __name__ == "__main__": - # For standalone testing - run_dispatcher() diff --git a/rapidfireai/evals/dispatcher/gunicorn.conf.py b/rapidfireai/evals/dispatcher/gunicorn.conf.py deleted file mode 100644 index b233bd62..00000000 --- a/rapidfireai/evals/dispatcher/gunicorn.conf.py +++ /dev/null @@ -1,25 +0,0 @@ -"""Gunicorn configuration file for the RapidFire dispatcher""" - -from rapidfireai.evals.db import RFDatabase -from rapidfireai.utils.constants import DispatcherConfig - -# Other Gunicorn settings... -bind = f"{DispatcherConfig.HOST}:{DispatcherConfig.PORT}" -workers = 1 # Single worker for Colab/single-user environments to save memory - -wsgi_app = "rapidfireai.evals.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 = RFDatabase() - rf_db.create_tables() - print("Database tables initialized successfully") - except Exception as e: - print(f"Error initializing database tables: {e}") - raise diff --git a/rapidfireai/evals/metrics/aggregator.py b/rapidfireai/evals/metrics/aggregator.py index f1d82d9e..02013e9b 100644 --- a/rapidfireai/evals/metrics/aggregator.py +++ b/rapidfireai/evals/metrics/aggregator.py @@ -1,6 +1,7 @@ from collections import defaultdict from collections.abc import Callable from typing import Any + import ray from .online_strategies import ( diff --git a/rapidfireai/evals/metrics/online_strategies.py b/rapidfireai/evals/metrics/online_strategies.py index e1a3dda9..996d63b3 100644 --- a/rapidfireai/evals/metrics/online_strategies.py +++ b/rapidfireai/evals/metrics/online_strategies.py @@ -3,9 +3,8 @@ from collections.abc import Callable from typing import Any -from typing_extensions import override - from scipy.stats import norm +from typing_extensions import override class OnlineAggregationStrategy(ABC): diff --git a/rapidfireai/evals/rag/prompt_manager.py b/rapidfireai/evals/rag/prompt_manager.py index 882853c9..525745b8 100644 --- a/rapidfireai/evals/rag/prompt_manager.py +++ b/rapidfireai/evals/rag/prompt_manager.py @@ -1,3 +1,7 @@ +import asyncio +import concurrent.futures +import hashlib +import json from copy import deepcopy from typing import Any @@ -206,6 +210,7 @@ async def _fetch(instructions: str, query: str) -> str: try: loop = asyncio.get_running_loop() + def run_in_thread(): new_loop = asyncio.new_event_loop() asyncio.set_event_loop(new_loop) @@ -262,9 +267,9 @@ def __getstate__(self): """ state = self.__dict__.copy() # Remove unpicklable objects - state['embedding_function'] = None - state['example_selector'] = None - state['fewshot_generator'] = None + state["embedding_function"] = None + state["example_selector"] = None + state["fewshot_generator"] = None return state def __setstate__(self, state): @@ -295,12 +300,9 @@ def get_hash(self) -> str: "example_selector_cls": self.example_selector_cls.__name__ if self.example_selector_cls else None, "num_examples": len(self.examples) if self.examples else 0, # Hash the examples themselves to detect changes - "examples_hash": hashlib.sha256( - json.dumps(self.examples, sort_keys=True).encode() - ).hexdigest() + "examples_hash": hashlib.sha256(json.dumps(self.examples, sort_keys=True).encode()).hexdigest() if self.examples else None, } prompt_json = json.dumps(prompt_dict, sort_keys=True) return hashlib.sha256(prompt_json.encode()).hexdigest() - diff --git a/rapidfireai/evals/rag/rag_pipeline.py b/rapidfireai/evals/rag/rag_pipeline.py index 37d9aaac..d0abd20c 100644 --- a/rapidfireai/evals/rag/rag_pipeline.py +++ b/rapidfireai/evals/rag/rag_pipeline.py @@ -3,6 +3,7 @@ """ +import builtins import copy import warnings from collections.abc import Callable @@ -13,7 +14,7 @@ import mlflow from mlflow.entities import SpanType -from rapidfireai.evals.utils.constants import SEARCH_DEFAULTS, VALID_SEARCH_TYPES +from rapidfireai.utils.constants import SEARCH_DEFAULTS, VALID_SEARCH_TYPES import faiss from pinecone import Pinecone, ServerlessSpec, PodSpec, ByocSpec @@ -31,6 +32,7 @@ from langchain_community.cross_encoders import HuggingFaceCrossEncoder from langchain_classic.retrievers.document_compressors import CrossEncoderReranker + class LangChainRagSpec: """ RAG (Retrieval-Augmented Generation) implementation using LangChain. @@ -71,7 +73,7 @@ def __init__( search_cfg: Optional[dict[str, Any]] | None = None, reranker_cfg: Optional[dict[str, Any]] | None = None, enable_gpu_search: bool = False, - document_template: Optional[Callable[[Document], str]] | None = None, + document_template: Callable[[Document], str] | None = None, ) -> None: """ Initialize the RAG specification with LangChain components. @@ -160,14 +162,14 @@ def __init__( cfg = dict(reranker_cfg) cls_value = cfg.pop("class", None) if isinstance(cls_value, str): - from rapidfireai.evals.utils.constants import RERANKER_CLASS_REGISTRY + from rapidfireai.utils.constants import RERANKER_CLASS_REGISTRY resolved = RERANKER_CLASS_REGISTRY.get(cls_value) if resolved is None: raise ValueError( f"Unknown reranker class '{cls_value}'. " f"Supported classes: {sorted(RERANKER_CLASS_REGISTRY.keys())}. " f"To add a new reranker, register it in RERANKER_CLASS_REGISTRY " - f"in rapidfireai/evals/utils/constants.py." + f"in rapidfireai/utils/constants.py." ) self.reranker_cls = resolved else: @@ -194,17 +196,97 @@ def __init__( def default_template(doc: Document) -> str: """ Default document formatting template. - + Args: doc: A langchain Document to format. - + Returns: Formatted string with metadata and content. """ metadata = "; ".join([f"{k}: {v}" for k, v in doc.metadata.items()]) return f"{metadata}:\n{doc.page_content}" - def _load_documents(self) -> list[Document]: + def build_index(self) -> None: + """ + Create the embedding instance and ensure a retriever is available. + + - If a retriever was provided at init, it is used as-is (embedding is still + created for any callers that need it). + - If a vector_store was provided but no retriever, a retriever is created from + it via as_retriever(search_type=..., search_kwargs=...). + - If neither retriever nor vector_store was provided, a FAISS vector store is + created, documents are loaded (and optionally split if text_splitter is set), + added to FAISS, and a retriever is created from the vector store. + + FAISS index type when building: + - enable_gpu_search=True: IndexFlatL2 on GPU (exact L2 search). + - enable_gpu_search=False: IndexHNSWFlat on CPU (approximate HNSW search). + + Raises: + Exception: If embedding instantiation fails. + ImportError: If faiss (or faiss-gpu) is not available when building index. + """ + # Create embedding instance with provided configuration + self.embedding = self.embedding_cls(**self.embedding_kwargs) + + if self.reranker_cls: + if self.reranker_cls is CrossEncoderReranker: + hf_model_name = self.reranker_kwargs.pop("model_name", "cross-encoder/ms-marco-MiniLM-L6-v2") + hf_model_kwargs = self.reranker_kwargs.pop("model_kwargs", {}) + + self.reranker = self.reranker_cls( + model=HuggingFaceCrossEncoder(model_name=hf_model_name, model_kwargs=hf_model_kwargs), + **self.reranker_kwargs, + ) + else: + self.reranker = self.reranker_cls(**self.reranker_kwargs) + + # Initialize vector store and retriever based on provided parameters + if not self.retriever and not self.vector_store: + try: + if self.enable_gpu_search: + # Use GPU-accelerated FAISS vector store with exact search (L2 distance) by default + # FAISS vector store will be built (adding documents) in _build_vector_store() method + self.vector_store = FAISS( + embedding_function=self.embedding, + index=faiss.IndexFlatL2(len(self.embedding.embed_query("RapidFire AI is awesome!"))), + docstore=InMemoryDocstore(), + index_to_docstore_id={}, + ) + + else: + # Use CPU-based Implementation of FAISS with approximate search HNSW + # TODO: move these to constants.py + M = 16 # good default value: controls the number of bidirectional connections of each node + ef_construction = 64 # 4-8x of M: Size of dynamic candidate list during construction + ef_search = 32 # 2-4x of M: Size of dynamic candidate list during search + + hnsw_index = faiss.IndexHNSWFlat(len(self.embedding.embed_query("RapidFire AI is awesome!")), M) + hnsw_index.hnsw.efConstruction = ef_construction + hnsw_index.hnsw.efSearch = ef_search + + self.vector_store = FAISS( + embedding_function=self.embedding, + index=hnsw_index, + docstore=InMemoryDocstore(), + index_to_docstore_id={}, + ) + + except ImportError as e: + raise ImportError( + "FAISS is required for GPU similarity search. Install it with: pip install faiss-gpu" + ) from e + + self._build_vector_store() + self.retriever = self.vector_store.as_retriever( + search_type=self.search_type, search_kwargs=self.search_kwargs + ) + elif not self.retriever: + self.retriever = self.vector_store.as_retriever( + search_type=self.search_type, search_kwargs=self.search_kwargs + ) + + def _load_documents(self) -> builtins.list[Document]: """ Load documents using the configured document loader. @@ -213,7 +295,7 @@ def _load_documents(self) -> list[Document]: """ return self.document_loader.load() if self.document_loader is not None else [] - def _split_documents(self, documents: list[Document]) -> list[Document]: + def _split_documents(self, documents: builtins.list[Document]) -> builtins.list[Document]: """ Split documents into smaller chunks using the configured text splitter. @@ -361,6 +443,31 @@ def _build_vector_store(self) -> None: for i in range(0, len(documents), batch_size): self.vector_store.add_documents(documents=documents[i: i + batch_size]) + def _rerank_docs( + self, batch_queries: builtins.list[str], batch_docs: builtins.list[builtins.list[Document]] + ) -> builtins.list[builtins.list[Document]]: + """ + Optionally rerank batch documents using the configured BaseDocumentCompressor. + + The reranker (BaseDocumentCompressor) is applied to each query's document list + individually using the compress_documents() method, which requires both the + query and documents as input. + + Args: + batch_queries: A list of query strings corresponding to each document list. + batch_docs: A batch of document lists where each inner list contains + documents for a single query. + + Returns: + List[List[Document]]: The batch of documents, reranked if a reranker + (BaseDocumentCompressor) is configured, otherwise + returned as-is. Maintains the same structure as input. + """ + if self.reranker: + return [self.reranker.compress_documents(docs, query) + for query, docs in zip(batch_queries, batch_docs)] + return batch_docs + def build_pipeline(self) -> None: """ Create the embedding instance and ensure a retriever is available. @@ -601,4 +708,4 @@ def _default(obj): return str(obj) rag_json = json.dumps(rag_dict, sort_keys=True, default=_default) - return hashlib.sha256(rag_json.encode()).hexdigest() \ No newline at end of file + return hashlib.sha256(rag_json.encode()).hexdigest() diff --git a/rapidfireai/evals/scheduling/controller.py b/rapidfireai/evals/scheduling/controller.py index 9fda78a3..b04dcff5 100644 --- a/rapidfireai/evals/scheduling/controller.py +++ b/rapidfireai/evals/scheduling/controller.py @@ -1,35 +1,32 @@ import hashlib import json import time -from collections.abc import Callable from typing import Any + import ray -from rapidfireai.utils.constants import ColabConfig, RF_EXPERIMENT_PATH +from rapidfireai.automl import RFGridSearch, RFRandomSearch, get_flattened_config_leaf, get_runs +from rapidfireai.db import RfDb from rapidfireai.evals.actors.doc_actor import DocProcessingActor from rapidfireai.evals.actors.inference_engines import InferenceEngine from rapidfireai.evals.actors.query_actor import QueryProcessingActor from rapidfireai.evals.data.dataset import DataLoader -from rapidfireai.evals.db import RFDatabase from rapidfireai.evals.metrics.aggregator import Aggregator from rapidfireai.evals.scheduling.interactive_control import InteractiveControlHandler from rapidfireai.evals.scheduling.pipeline_scheduler import PipelineScheduler from rapidfireai.evals.scheduling.scheduler import Scheduler -from rapidfireai.automl import ModelConfig, RFvLLMModelConfig -from rapidfireai.evals.utils.constants import ( +from rapidfireai.evals.utils.progress_display import ContextBuildingDisplay, PipelineProgressDisplay +from rapidfireai.utils.constants import ( + ContextStatus, NUM_CPUS_PER_DOC_ACTOR, NUM_QUERY_PROCESSING_ACTORS, SEARCH_TYPE_KEYS, - ContextStatus, - ExperimentStatus, PipelineStatus, + RF_EXPERIMENT_PATH, TaskStatus, ) -from rapidfireai.evals.utils.logger import RFLogger -from rapidfireai.evals.utils.progress_display import ContextBuildingDisplay, PipelineProgressDisplay -from rapidfireai.automl import RFGridSearch, RFRandomSearch -from rapidfireai.automl import get_runs, get_flattened_config_leaf -from rapidfireai.evals.utils.serialize import extract_pipeline_config_json, extract_pipeline_display_metadata +from rapidfireai.utils.logging import RFLogger +from rapidfireai.utils.serialize import extract_pipeline_config_json class Controller: """ @@ -270,7 +267,7 @@ def _collect_unique_contexts( def _setup_context_generators( self, config_leaves: Any, - db: RFDatabase, + db: RfDb, ) -> None: """ Setup RAG contexts: build if needed and cache in Ray object store. @@ -353,7 +350,7 @@ def _setup_context_generators( def build_rag_components( self, contexts_to_build: list[dict], - db: RFDatabase, + db: RfDb, ) -> None: """ Build multiple RAG components in parallel using DocProcessingActors. @@ -501,13 +498,13 @@ def build_rag_components( # HALT: Context creation is critical - stop the entire experiment context_display.stop() error_message = ( - f"\n{'='*80}\n" + f"\n{'=' * 80}\n" f"โŒ CRITICAL ERROR: RAG Source Preprocessing Failed\n" - f"{'='*80}\n" + f"{'=' * 80}\n" f"RAG Source ID: {context_id}\n" f"Context Hash: {context_hash[:16]}...\n" f"Error: {str(e)}\n" - f"{'='*80}\n" + f"{'=' * 80}\n" f"\nThe experiment has been halted. Please fix the error and try again.\n" ) print(error_message) @@ -569,7 +566,7 @@ def create_query_actors( def _register_pipelines( self, config_leaves: list[Any], - db: RFDatabase, + db: RfDb, ) -> tuple[list[int], dict[int, dict]]: """ Register pipelines in database. @@ -641,7 +638,9 @@ def _register_pipelines( if hasattr(pipeline, "rag") and pipeline.rag: if hasattr(pipeline.rag, "search_type"): - self.metric_manager.log_param(metric_run_id, "rag_search_type", str(pipeline.rag.search_type)) + self.metric_manager.log_param( + metric_run_id, "rag_search_type", str(pipeline.rag.search_type) + ) if hasattr(pipeline.rag, "search_kwargs") and pipeline.rag.search_kwargs: k = pipeline.rag.search_kwargs.get("k") if k is not None: @@ -650,7 +649,12 @@ def _register_pipelines( # Extract sampling params if hasattr(pipeline, "sampling_params") and pipeline.sampling_params: import json - sampling_str = json.dumps(pipeline.sampling_params) if isinstance(pipeline.sampling_params, dict) else str(pipeline.sampling_params) + + sampling_str = ( + json.dumps(pipeline.sampling_params) + if isinstance(pipeline.sampling_params, dict) + else str(pipeline.sampling_params) + ) self.metric_manager.log_param(metric_run_id, "sampling_params", sampling_str) self.logger.debug(f"Created Metrics run {metric_run_id} for pipeline {pipeline_id}") @@ -668,7 +672,7 @@ def _compute_final_metrics_for_pipelines( pipeline_id_to_config: dict[int, dict], pipeline_aggregators: dict[int, Aggregator], pipeline_results: dict[int, dict], - db: RFDatabase, + db: RfDb, progress_display=None, pipeline_id_to_info: dict[int, dict] = None, total_dataset_size: int = None, @@ -733,8 +737,7 @@ def _compute_final_metrics_for_pipelines( # Add confidence intervals to final metrics before storing samples_processed = sum( - m.get("value", 0) - for m in pipeline_results[pipeline_id]["metrics"].get("Samples Processed", [{}]) + m.get("value", 0) for m in pipeline_results[pipeline_id]["metrics"].get("Samples Processed", [{}]) ) if aggregator.online_strategy and samples_processed > 0: cumulative_metrics = aggregator.online_strategy.add_confidence_interval_info( @@ -799,11 +802,7 @@ def _compute_final_metrics_for_pipelines( else: display_metrics[metric_name] = {"value": metric_data} - progress_display.update_pipeline( - pipeline_id, - status="COMPLETED", - metrics=display_metrics - ) + progress_display.update_pipeline(pipeline_id, status="COMPLETED", metrics=display_metrics) # Log final metrics to MetricLogger if self.metric_manager: @@ -813,13 +812,19 @@ def _compute_final_metrics_for_pipelines( if metric_run_id: total_samples = pipeline.get("total_samples_processed", 0) if total_dataset_size and total_dataset_size > 0: - percentage_processed = (total_samples / total_dataset_size * 100) + percentage_processed = total_samples / total_dataset_size * 100 else: percentage_processed = 100 # Assume complete if dataset size unknown step = int(percentage_processed) for metric_name, metric_data in ordered_metrics.items(): - if metric_name in ["run_id", "model_name", "Samples Processed", "Processing Time", "Samples Per Second"]: + if metric_name in [ + "run_id", + "model_name", + "Samples Processed", + "Processing Time", + "Samples Per Second", + ]: continue if isinstance(metric_data, dict): @@ -832,7 +837,9 @@ def _compute_final_metrics_for_pipelines( # Log main metric value if isinstance(metric_value, (int, float)): try: - self.metric_manager.log_metric(metric_run_id, metric_name, float(metric_value), step=step) + self.metric_manager.log_metric( + metric_run_id, metric_name, float(metric_value), step=step + ) except Exception as e: self.logger.debug(f"Failed to log final metric {metric_name} to MetricLogger: {e}") @@ -887,7 +894,7 @@ def run_multi_pipeline_inference( Dict mapping pipeline_id to (aggregated_results, cumulative_metrics) tuple """ # Initialize database - db = RFDatabase() + db = RfDb() # PHASE 1: Shard the dataset shards = self.dataloader.get_shards_from_data(dataset, num_shards) @@ -1055,6 +1062,7 @@ def run_multi_pipeline_inference( for pipeline_id, pipeline_config in pipeline_id_to_config.items(): from rapidfireai.automl import RFOpenAIAPIModelConfig + pipeline = pipeline_config["pipeline"] if hasattr(pipeline, "model_config") and isinstance(pipeline, RFOpenAIAPIModelConfig): has_openai_pipeline = True @@ -1073,16 +1081,17 @@ def run_multi_pipeline_inference( "tpm": pipeline.tpm_limit, } max_completion_tokens_by_model[model_name] = model_config.get("max_completion_tokens", 150) - elif (model_rate_limits[model_name]["rpm"] != pipeline.rpm_limit or - model_rate_limits[model_name]["tpm"] != pipeline.tpm_limit): + elif ( + model_rate_limits[model_name]["rpm"] != pipeline.rpm_limit + or model_rate_limits[model_name]["tpm"] != pipeline.tpm_limit + ): self.logger.warning( f"Model {model_name} has inconsistent rate limits across pipelines. " f"Using first encountered values: {model_rate_limits[model_name]}" ) pipeline_to_max_completion_tokens[pipeline_id] = model_config.get( - "max_completion_tokens", - max_completion_tokens_by_model.get(model_name, 150) + "max_completion_tokens", max_completion_tokens_by_model.get(model_name, 150) ) pipeline_to_rate_limiter[pipeline_id] = None @@ -1107,14 +1116,12 @@ def run_multi_pipeline_inference( pipeline = pipeline_config["pipeline"] model_name = pipeline.model_config.get("model", "gpt-3.5-turbo") pipeline_to_max_completion_tokens[pipeline_id] = max_completion_tokens_by_model.get( - model_name, - pipeline.model_config.get("max_completion_tokens", 150) + model_name, pipeline.model_config.get("max_completion_tokens", 150) ) - limits_summary = ", ".join([ - f"{model}: {limits['rpm']} RPM, {limits['tpm']} TPM" - for model, limits in model_rate_limits.items() - ]) + limits_summary = ", ".join( + [f"{model}: {limits['rpm']} RPM, {limits['tpm']} TPM" for model, limits in model_rate_limits.items()] + ) self.logger.info( f"Created experiment-wide rate limiter actor for {len(pipeline_to_rate_limiter)} OpenAI pipeline(s) " f"with per-model limits ({limits_summary})" @@ -1189,9 +1196,7 @@ def run_multi_pipeline_inference( # Mark as completed (metrics will be finalized in Phase 8) db.set_pipeline_status(pipeline_id, PipelineStatus.COMPLETED) progress_display.update_pipeline(pipeline_id, status="COMPLETED") - self.logger.info( - f"Pipeline {pipeline_id} completed all {num_shards} shards" - ) + self.logger.info(f"Pipeline {pipeline_id} completed all {num_shards} shards") # Compute current metrics with confidence intervals confidence_value = None @@ -1247,12 +1252,24 @@ def run_multi_pipeline_inference( if self.metric_manager and metric_run_id: try: - actual_samples_processed = pipeline.get("total_samples_processed", samples_processed) - percentage_processed = (actual_samples_processed / total_dataset_size * 100) if total_dataset_size > 0 else 0 + actual_samples_processed = pipeline.get( + "total_samples_processed", samples_processed + ) + percentage_processed = ( + (actual_samples_processed / total_dataset_size * 100) + if total_dataset_size > 0 + else 0 + ) step = int(percentage_processed) for metric_name, metric_data in metrics_with_ci.items(): - if metric_name in ["run_id", "model_name", "Samples Processed", "Processing Time", "Samples Per Second"]: + if metric_name in [ + "run_id", + "model_name", + "Samples Processed", + "Processing Time", + "Samples Per Second", + ]: continue if isinstance(metric_data, dict): @@ -1265,9 +1282,13 @@ def run_multi_pipeline_inference( # Log main metric value if isinstance(metric_value, (int, float)): try: - self.metric_manager.log_metric(metric_run_id, metric_name, float(metric_value), step=step) + self.metric_manager.log_metric( + metric_run_id, metric_name, float(metric_value), step=step + ) except Exception as e: - self.logger.debug(f"Failed to log metric {metric_name} to MetricLogger: {e}") + self.logger.debug( + f"Failed to log metric {metric_name} to MetricLogger: {e}" + ) # Log confidence_interval if available if confidence_interval is not None and isinstance(confidence_interval, (int, float)): @@ -1280,7 +1301,12 @@ def run_multi_pipeline_inference( throughput_value = display_metrics["Throughput"]["value"] if isinstance(throughput_value, (int, float)): try: - self.metric_manager.log_metric(metric_run_id, "Throughput", float(throughput_value), step=step) + self.metric_manager.log_metric( + metric_run_id, + "Throughput", + float(throughput_value), + step=step, + ) except Exception as e: self.logger.debug(f"Failed to log Throughput to MetricLogger: {e}") except Exception as e: @@ -1317,14 +1343,18 @@ def run_multi_pipeline_inference( # Display error in notebook (but don't halt the experiment) pipeline_config = pipeline_id_to_config.get(pipeline_id) - pipeline_name = pipeline_config.get("pipeline_name", f"Pipeline {pipeline_id}") if pipeline_config else f"Pipeline {pipeline_id}" + pipeline_name = ( + pipeline_config.get("pipeline_name", f"Pipeline {pipeline_id}") + if pipeline_config + else f"Pipeline {pipeline_id}" + ) error_display = ( - f"\n{'='*80}\n" + f"\n{'=' * 80}\n" f"โš ๏ธ Run {pipeline_id} ({pipeline_name}) FAILED\n" - f"{'='*80}\n" + f"{'=' * 80}\n" f"Shard: {shard_id + 1}/{num_shards}\n" f"Error: {error_msg}\n" - f"{'='*80}\n" + f"{'=' * 80}\n" f"This run has been marked as FAILED. The experiment will continue with other runs.\n" ) print(error_display) @@ -1415,7 +1445,7 @@ def run_multi_pipeline_inference( db.set_pipeline_status(pipeline_id, PipelineStatus.ONGOING) # Get shard data and split into batches - batch_size = pipeline_config["batch_size"]#TODO: set default batch size + batch_size = pipeline_config["batch_size"] # TODO: set default batch size shard_data = shards[shard_id] batches = self.dataloader.get_batches(shard_data, batch_size) @@ -1588,7 +1618,6 @@ def run_multi_pipeline_inference( total_dataset_size=total_dataset_size, ) - # Cleanup actors for actor in query_actors: ray.kill(actor) diff --git a/rapidfireai/evals/scheduling/interactive_control.py b/rapidfireai/evals/scheduling/interactive_control.py index 30edfa97..4b78217d 100644 --- a/rapidfireai/evals/scheduling/interactive_control.py +++ b/rapidfireai/evals/scheduling/interactive_control.py @@ -12,15 +12,15 @@ import time from rapidfireai.evals.actors.rate_limiter_actor import RateLimiterActor -from rapidfireai.evals.utils.constants import RERANKER_CLASS_REGISTRY, SEARCH_DEFAULTS, SEARCH_TYPE_KEYS - -from rapidfireai.evals.db import RFDatabase +from rapidfireai.automl import RFOpenAIAPIModelConfig, RFvLLMModelConfig +from rapidfireai.db import RfDb from rapidfireai.evals.metrics.aggregator import Aggregator from rapidfireai.evals.scheduling.pipeline_scheduler import PipelineScheduler -from rapidfireai.automl import RFOpenAIAPIModelConfig, RFvLLMModelConfig -from rapidfireai.evals.utils.constants import ICOperation, ICStatus, PipelineStatus -from rapidfireai.evals.utils.logger import RFLogger -from rapidfireai.evals.utils.serialize import extract_pipeline_display_metadata +from rapidfireai.utils.constants import ( + ICOperation, ICStatus, PipelineStatus, + RERANKER_CLASS_REGISTRY, SEARCH_DEFAULTS, SEARCH_TYPE_KEYS, +) +from rapidfireai.utils.logging import RFLogger class InteractiveControlHandler: @@ -54,7 +54,7 @@ def __init__(self, experiment_name: str, experiment_path: str, context_cache: di def check_and_process_requests( self, scheduler: PipelineScheduler, - db: RFDatabase, + db: RfDb, num_shards: int, dataset, pipeline_aggregators: dict, @@ -81,14 +81,14 @@ def check_and_process_requests( online_strategy_kwargs: Optional online aggregation strategy parameters progress_display: Optional PipelineProgressDisplay instance for updating UI """ - pending_ops = db.get_pending_ic_operations() + pending_ops = db.get_pending_ic_operations(target_type="pipeline") if not pending_ops: return for op in pending_ops: ic_id = op["ic_id"] operation = op["operation"] - pipeline_id = op["pipeline_id"] + pipeline_id = op["target_id"] # Unified schema uses target_id for both runs and pipelines try: # Mark as processing @@ -105,7 +105,7 @@ def check_and_process_requests( elif operation == ICOperation.CLONE.value: self._handle_clone( - op["request_data"], + op["config_data"], # Unified schema uses config_data scheduler, db, num_shards, @@ -137,7 +137,7 @@ def check_and_process_requests( time.sleep(0.5) def _handle_stop( - self, pipeline_id: int, scheduler: PipelineScheduler, db: RFDatabase, progress_display=None + self, pipeline_id: int, scheduler: PipelineScheduler, db: RfDb, progress_display=None ) -> None: """ Stop a pipeline (remove from scheduling, save progress). @@ -161,7 +161,7 @@ def _handle_stop( self.logger.info(f"Stopped pipeline {pipeline_id} at {shards_completed} shards completed") def _handle_resume( - self, pipeline_id: int, scheduler: PipelineScheduler, db: RFDatabase, num_shards: int, progress_display=None + self, pipeline_id: int, scheduler: PipelineScheduler, db: RfDb, num_shards: int, progress_display=None ) -> None: """ Resume a stopped pipeline (re-add to scheduler with saved progress). @@ -201,7 +201,7 @@ def _handle_delete( self, pipeline_id: int, scheduler: PipelineScheduler, - db: RFDatabase, + db: RfDb, pipeline_results: dict, progress_display=None, ) -> None: @@ -242,9 +242,9 @@ def _handle_delete( def _handle_clone( self, - request_data: str, + request_data: dict | str, scheduler: PipelineScheduler, - db: RFDatabase, + db: RfDb, num_shards: int, dataset, pipeline_aggregators: dict, @@ -261,7 +261,7 @@ def _handle_clone( Only the JSON-editable parameters (model config, sampling params, etc.) are modified. Args: - request_data: JSON string with {"parent_pipeline_id": int, "config_json": {...}} + request_data: Dict or JSON string with {"parent_pipeline_id": int, "config_json": {...}} scheduler: PipelineScheduler instance db: Database instance num_shards: Total number of shards @@ -275,8 +275,11 @@ def _handle_clone( Returns: pipeline_id of the newly created pipeline """ - # Parse request data - data = json.loads(request_data) + # Parse request data (handle both dict and JSON string) + if isinstance(request_data, str): + data = json.loads(request_data) + else: + data = request_data parent_pipeline_id = data["parent_pipeline_id"] edited_json = data["config_json"] @@ -359,14 +362,14 @@ def _handle_clone( f"Unknown reranker class '{cls_value}'. " f"Supported classes: {sorted(RERANKER_CLASS_REGISTRY.keys())}. " f"To use a custom class, add it to RERANKER_CLASS_REGISTRY in " - f"rapidfireai/evals/utils/constants.py." + f"rapidfireai/utils/constants.py." ) rag.reranker_cls = resolved if rag.reranker_kwargs is None: rag.reranker_kwargs = {} rag.reranker_kwargs.update({k: v for k, v in reranker_cfg.items() if k != "class"}) - self.logger.info(f"RAG config updated successfully") + self.logger.info("RAG config updated successfully") # Extract pipeline type from edited JSON (or inherit from parent) pipeline_type = edited_json.get("pipeline_type") @@ -453,7 +456,7 @@ def _handle_clone( # Merge parent online_strategy_kwargs with user edits pipeline_config_dict["online_strategy_kwargs"] = { **parent_online_strategy, - **edited_json.get("online_strategy_kwargs", {}) + **edited_json.get("online_strategy_kwargs", {}), } elif "online_strategy_kwargs" in edited_json: # No parent strategy, use user's strategy as-is diff --git a/rapidfireai/evals/scheduling/scheduler.py b/rapidfireai/evals/scheduling/scheduler.py index 9f0879fe..ee3bb251 100644 --- a/rapidfireai/evals/scheduling/scheduler.py +++ b/rapidfireai/evals/scheduling/scheduler.py @@ -6,6 +6,7 @@ from abc import ABC, abstractmethod from collections.abc import Callable from typing import Any + import ray diff --git a/rapidfireai/evals/utils/__init__.py b/rapidfireai/evals/utils/__init__.py index 8b39cc41..d00e5ff9 100644 --- a/rapidfireai/evals/utils/__init__.py +++ b/rapidfireai/evals/utils/__init__.py @@ -1,6 +1,6 @@ """Utility functions and helpers.""" -from .constants import get_dispatcher_headers, get_dispatcher_url +from rapidfireai.utils.constants import get_dispatcher_headers, get_dispatcher_url __all__ = [ "get_dispatcher_url", diff --git a/rapidfireai/evals/utils/constants.py b/rapidfireai/evals/utils/constants.py deleted file mode 100644 index 281a1fa2..00000000 --- a/rapidfireai/evals/utils/constants.py +++ /dev/null @@ -1,171 +0,0 @@ -import os -from enum import Enum -from rapidfireai.utils.colab import get_colab_auth_token -from rapidfireai.utils.constants import DispatcherConfig, ColabConfig, RF_DB_PATH, ExperimentStatus - -# --------------------------------------------------------------------------- -# Reranker class registry -# --------------------------------------------------------------------------- -# Maps the string that appears in the clone-modify dialog (the class __qualname__) -# back to the live class object so we can restore it after round-tripping through -# JSON. Add any new reranker class here to make it available in the UI. -# --------------------------------------------------------------------------- -def _build_reranker_registry() -> dict[str, type]: - registry: dict[str, type] = {} - try: - from langchain_classic.retrievers.document_compressors import CrossEncoderReranker - registry[CrossEncoderReranker.__qualname__] = CrossEncoderReranker - except ImportError: - pass - try: - from langchain.retrievers.document_compressors import LLMChainExtractor - registry[LLMChainExtractor.__qualname__] = LLMChainExtractor - except ImportError: - pass - try: - from langchain_cohere import CohereRerank - registry[CohereRerank.__qualname__] = CohereRerank - except ImportError: - pass - return registry - -RERANKER_CLASS_REGISTRY: dict[str, type] = _build_reranker_registry() - -# Actor Constants -NUM_QUERY_PROCESSING_ACTORS = 4 -NUM_CPUS_PER_DOC_ACTOR = 2 if os.cpu_count() > 2 else 1 - -# RAG search type defaults โ€” keys are type-specific; only include kwargs relevant -# to each search type so irrelevant params are never forwarded to LangChain. -VALID_SEARCH_TYPES = {"similarity", "similarity_score_threshold", "mmr"} -SEARCH_DEFAULTS: dict[str, dict] = { - "similarity": {"k": 5, "filter": None}, - "similarity_score_threshold": {"k": 5, "filter": None, "score_threshold": 0.5}, - "mmr": {"k": 5, "filter": None, "fetch_k": 20, "lambda_mult": 0.5}, -} -# Keys shown/accepted per search type (used by serialize and interactive_control) -SEARCH_TYPE_KEYS: dict[str, set] = {search_type: set(defaults.keys()) for search_type, defaults in SEARCH_DEFAULTS.items()} - -# Rate Limiting Constants -# Maximum number of retries for rate-limited API calls -MAX_RATE_LIMIT_RETRIES = 5 -# Base wait time for exponential backoff (seconds) -RATE_LIMIT_BACKOFF_BASE = 2 - -def get_dispatcher_url() -> str: - """ - Auto-detect dispatcher URL based on environment. - - Returns: - - In Google Colab: Uses Colab's kernel proxy URL (e.g., https://xxx-8851-xxx.ngrok-free.app) - - In Jupyter/Local: Uses localhost URL (http://127.0.0.1:8851) - """ - if ColabConfig.ON_COLAB: - try: - from google.colab.output import eval_js - - # Get the Colab proxy URL for the dispatcher port - proxy_url = eval_js(f"google.colab.kernel.proxyPort({DispatcherConfig.PORT})") - print(f"๐ŸŒ Google Colab detected. Dispatcher URL: {proxy_url}") - return proxy_url - except Exception as e: - print(f"โš ๏ธ Colab detected but failed to get proxy URL: {e}") - # Fall back to localhost - return DispatcherConfig.URL - else: - # Running in Jupyter, local Python, or other environment - return DispatcherConfig.URL - - - -def get_dispatcher_headers() -> dict[str, str]: - """ - Get the HTTP headers needed for dispatcher API requests. - - Returns: - Dictionary with required headers, including Authorization header in Colab - """ - headers = {"Content-Type": "application/json"} - - auth_token = get_colab_auth_token() - if auth_token: - headers["Authorization"] = f"Bearer {auth_token}" - - return headers - - -# TODO: Merge multiple Statuses into a single Status enum - -# Note: ExperimentStatus is imported from rapidfireai.utils.constants (shared) - - -class ContextStatus(str, Enum): - """Status values for RAG contexts.""" - - NEW = "new" - ONGOING = "ongoing" - DELETED = "deleted" - FAILED = "failed" - - -class PipelineStatus(str, Enum): - """Status values for pipelines.""" - - NEW = "new" - ONGOING = "ongoing" - COMPLETED = "completed" - STOPPED = "stopped" - DELETED = "deleted" - FAILED = "failed" - - -class TaskStatus(str, Enum): - """Status values for actor tasks.""" - - SCHEDULED = "scheduled" - IN_PROGRESS = "in_progress" - COMPLETED = "completed" - FAILED = "failed" - - -class ICOperation(str, Enum): - """Interactive Control operation types.""" - - STOP = "stop" - RESUME = "resume" - DELETE = "delete" - CLONE = "clone" - - -class ICStatus(str, Enum): - """Status values for Interactive Control operations.""" - - PENDING = "pending" - PROCESSING = "processing" - COMPLETED = "completed" - FAILED = "failed" - - -# Database Constants -class DBConfig: - """Class to manage the database configuration for SQLite""" - - # Use user's home directory for database path - - DB_PATH: str = os.path.join( - RF_DB_PATH, "rapidfire_evals.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 diff --git a/rapidfireai/evals/utils/experiment_utils.py b/rapidfireai/evals/utils/experiment_utils.py deleted file mode 100644 index 5b33c9c8..00000000 --- a/rapidfireai/evals/utils/experiment_utils.py +++ /dev/null @@ -1,257 +0,0 @@ -"""This module contains utility functions for the evals experiment.""" - -import os -import re -import warnings -from pathlib import Path - -from rapidfireai.utils.constants import RF_EXPERIMENT_PATH -from rapidfireai.evals.db.rf_db import RFDatabase -from rapidfireai.evals.utils.constants import ExperimentStatus -from rapidfireai.evals.utils.logger import RFLogger - - -class ExperimentUtils: - """Class to contain utility functions for the experiment.""" - - def __init__(self) -> None: - # initialize database handler - self.db = RFDatabase() - - def _disable_ml_warnings_display(self) -> None: - """Disable warnings""" - # Suppress warnings - warnings.filterwarnings("ignore", message=".*torch.cuda.amp.autocast.*") - warnings.filterwarnings("ignore", message=".*torch.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.*") - warnings.filterwarnings("ignore", message=".*Unable to register cuDNN factory.*") - warnings.filterwarnings("ignore", message=".*Unable to register cuBLAS factory.*") - warnings.filterwarnings("ignore", message=".*All log messages before absl::InitializeLog.*") - warnings.filterwarnings("ignore", message=".*resource_tracker: process died unexpectedly.*") - warnings.filterwarnings("ignore", message=".*computation placer already registered.") - warnings.filterwarnings("ignore", message=".*Rank 0 is connected to 0 peer ranks.*") - warnings.filterwarnings("ignore", message=".*No cudagraph will be used.*") - warnings.filterwarnings("ignore", module="multiprocessing.resource_tracker") - - 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. - - Args: - given_name: Desired experiment name - experiments_path: Path to experiments directory - - Returns: - Tuple of (experiment_id, experiment_name, log_messages) - """ - log_messages: list[str] = [] - - # disable warnings - self._disable_ml_warnings_display() - - # check if experiment is already running - running_experiment = None - try: - running_experiment = self.db.get_running_experiment() - except Exception: - 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) - - # Cancel any running tasks - msg = "Any running tasks have 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: - # Different experiment - end the previous one and create new - self.end_experiment(internal=True) - experiment_id, experiment_name = self._create_experiment(given_name, experiments_path) - msg = ( - f"The previously running experiment {running_experiment['experiment_name']} was forcibly ended." - f" Created a new experiment '{experiment_name}' with Experiment ID: {experiment_id}" - f" 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 = self._create_experiment(given_name, experiments_path) - msg = ( - "An experiment with the same name already exists." - f" Created a new experiment '{experiment_name}' with Experiment ID: {experiment_id}" - f" at {experiments_path}/{experiment_name}" - ) - print(msg) - log_messages.append(msg) - else: - # New experiment - experiment_id, experiment_name = self._create_experiment(given_name, experiments_path) - msg = ( - f"Experiment {experiment_name} created with Experiment ID: {experiment_id}" - f" at {experiments_path}/{experiment_name}" - ) - print(msg) - log_messages.append(msg) - - # Create experiment directory - try: - experiment_dir = Path(experiments_path) / experiment_name - experiment_dir.mkdir(parents=True, exist_ok=True) - except Exception as e: - raise Exception(f"Failed to create experiment directories: {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 Exception: - if not internal: - print("No experiment is currently running. Nothing to end.") - return - - # Check if there's actually a running experiment - if current_experiment is None: - if not internal: - print("No experiment is currently running. Nothing to end.") - return - - # create logger - experiment_name = current_experiment["experiment_name"] - logging_manager = RFLogger(experiment_name=experiment_name, experiment_path=RF_EXPERIMENT_PATH) - logger = logging_manager.get_logger("ExperimentUtils") - - # cancel current tasks if 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() - - # 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 - marks experiment as cancelled and resets pipeline/context states""" - # check if experiment is running - try: - current_experiment = self.db.get_running_experiment() - except Exception: - if not internal: - print("No experiment is currently running. Nothing to cancel.") - return - - # Check if there's actually a running experiment - if current_experiment is None: - if not internal: - print("No experiment is currently running. Nothing to cancel.") - return - - # create logger - experiment_name = current_experiment["experiment_name"] - logging_manager = RFLogger(experiment_name=experiment_name, experiment_path=RF_EXPERIMENT_PATH) - logger = logging_manager.get_logger("ExperimentUtils") - - try: - # Reset experiment states (mark pipelines/contexts/tasks as failed) - self.db.reset_experiment_states() - logger.info("Reset experiment states - marked ongoing pipelines, contexts, and tasks as failed") - - # Mark experiment as cancelled - self.db.set_experiment_status(current_experiment["experiment_id"], ExperimentStatus.CANCELLED) - - msg = "Experiment marked as cancelled. Ongoing pipelines, contexts, and tasks have been marked as failed." - if not internal: - print(msg) - logger.info(msg) - except Exception as e: - msg = f"Error cancelling experiment: {e}" - if not internal: - print(msg) - logger.error(msg) - - def _create_experiment(self, given_name: str, experiments_path: str) -> tuple[int, str]: - """ - Create new experiment - if given_name already exists, increment suffix and create new experiment. - - Args: - given_name: Desired experiment name - experiments_path: Path to experiments directory - - Returns: - Tuple of (experiment_id, experiment_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()) - - # Clear all tables except experiments table before creating new experiment - # This ensures a clean slate for the new experiment - self.db.reset_all_tables(experiments_table=False) - - # write new experiment details to database - experiment_id = self.db.create_experiment( - experiment_name=experiment_name, - num_actors=0, # Will be updated in run_evals - status=ExperimentStatus.RUNNING, - ) - return experiment_id, experiment_name - except Exception as e: - raise Exception(f"Error in _create_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 Exception("Could not generate unique experiment name") - - return new_name \ No newline at end of file diff --git a/rapidfireai/evals/utils/logger.py b/rapidfireai/evals/utils/logger.py deleted file mode 100644 index cb8adfeb..00000000 --- a/rapidfireai/evals/utils/logger.py +++ /dev/null @@ -1,119 +0,0 @@ -import logging -import os -from pathlib import Path - -from rapidfireai.utils.constants import RF_LOG_FILENAME, RF_LOG_PATH, RF_EXPERIMENT_PATH -from rapidfireai.utils.os_utils import mkdir_p - - -class SafeLoggerAdapter(logging.LoggerAdapter): - """Custom LoggerAdapter that prefixes messages with experiment and logger name.""" - def process(self, msg, kwargs): - # Prefix message with experiment and logger name - experiment = self.extra.get("experiment_name", "unknown") - log_name = self.extra.get("logger_name", "unknown") - return f"[{experiment}:{log_name}] {msg}", kwargs - - -class RFLogger: - _file_handler = None - _experiment_name = None - _experiment_path = None - - def __init__( - self, experiment_name: str = "unknown", experiment_path: str = RF_EXPERIMENT_PATH, level: str = "INFO" - ): - self._experiment_name = experiment_name - self._experiment_path = experiment_path - self.level = level.upper() - - # Suppress third-party library logs via environment variables - os.environ.setdefault("VLLM_LOGGING_LEVEL", "ERROR") - os.environ.setdefault("RAY_LOG_TO_STDERR", "0") - - # Check if we are in a Ray worker process - is_ray_worker = os.environ.get("RAY_WORKER_MODE") == "WORKER" - - # Only set up the file handler on the Controller/Experiment process - # Ray workers will forward their logs to the driver's logging system - if not is_ray_worker and RFLogger._file_handler is None: - log_dir = Path(RF_LOG_PATH) / self._experiment_name - try: - mkdir_p(log_dir.absolute()) - except (PermissionError, OSError) as e: - print(f"Error creating directory: {e}") - raise - - # Use standard format fields only - LoggerAdapter will prefix messages - log_format = "%(asctime)s | %(name)s | %(levelname)s | %(filename)s:%(lineno)d | %(message)s" - - # Set up the file handler - log_file_path = log_dir / RF_LOG_FILENAME - RFLogger._file_handler = logging.FileHandler(log_file_path) - RFLogger._file_handler.setFormatter(logging.Formatter(log_format, datefmt="%Y-%m-%d %H:%M:%S")) - RFLogger._file_handler.setLevel(self.level) - - # Get root logger - root_logger = logging.getLogger() - - # Remove all existing handlers to prevent console output - for handler in root_logger.handlers[:]: - root_logger.removeHandler(handler) - - # Add only our file handler to the root logger - root_logger.addHandler(RFLogger._file_handler) - root_logger.setLevel(self.level) - - # Add filter to suppress harmless asyncio cleanup errors - class AsyncioCleanupFilter(logging.Filter): - """Filter out harmless asyncio cleanup errors during shutdown.""" - def filter(self, record): - # Suppress TCPTransport cleanup errors - if "TCPTransport" in str(record.getMessage()) and "closed=True" in str(record.getMessage()): - return False - # Suppress "Task exception was never retrieved" for these specific cases - if "Task exception was never retrieved" in str(record.getMessage()): - # Only suppress if it's related to httpx/OpenAI client cleanup - if "AsyncClient.aclose" in str(record.getMessage()): - return False - return True - - # Add the filter to the root logger - root_logger.addFilter(AsyncioCleanupFilter()) - - # Also suppress asyncio logger specifically - asyncio_logger = logging.getLogger("asyncio") - asyncio_logger.addFilter(AsyncioCleanupFilter()) - asyncio_logger.setLevel(logging.CRITICAL) - - # Suppress third-party library logs more aggressively - third_party_loggers = [ - "ray", - "vllm", - "torch", - "transformers", - "datasets", - "huggingface_hub", - "langchain", - "langchain_core", - "langchain_community", - "openai", - "httpx", - "urllib3", - ] - for logger_name in third_party_loggers: - logging.getLogger(logger_name).setLevel(logging.CRITICAL) - logging.getLogger(logger_name).propagate = False - - def get_logger(self, logger_name: str = "unknown"): - logger = logging.getLogger(logger_name) - logger.setLevel(self.level) - - return SafeLoggerAdapter( - logger, - { - "experiment_name": self._experiment_name, - "logger_name": logger_name, - }, - ) - diff --git a/rapidfireai/evals/utils/notebook_ui.py b/rapidfireai/evals/utils/notebook_ui.py index 002512a4..d9ac275d 100644 --- a/rapidfireai/evals/utils/notebook_ui.py +++ b/rapidfireai/evals/utils/notebook_ui.py @@ -6,13 +6,19 @@ import uuid from IPython.display import HTML, display -from rapidfireai.utils.constants import DispatcherConfig, ColabConfig + +from rapidfireai.utils.constants import ColabConfig, DispatcherConfig class NotebookUI: """Notebook UI that works in VS Code""" - def __init__(self, dispatcher_url: str = "http://127.0.0.1:8851", refresh_rate_seconds: float = 3.0, auth_token: str | None = None): + def __init__( + self, + dispatcher_url: str = "http://127.0.0.1:8851", + refresh_rate_seconds: float = 3.0, + auth_token: str | None = None, + ): if ColabConfig.ON_COLAB: self.dispatcher_url = f"https://localhost:{DispatcherConfig.PORT}" else: diff --git a/rapidfireai/evals/utils/ratelimiter.py b/rapidfireai/evals/utils/ratelimiter.py index 75994a98..42cf8c15 100644 --- a/rapidfireai/evals/utils/ratelimiter.py +++ b/rapidfireai/evals/utils/ratelimiter.py @@ -62,21 +62,15 @@ def __init__( self.model_names = list(model_rate_limits.keys()) # Actual API limits per model (as specified by the API provider) - self.actual_rpm_limits: dict[str, int] = { - model: limits["rpm"] for model, limits in model_rate_limits.items() - } - self.actual_tpm_limits: dict[str, int] = { - model: limits["tpm"] for model, limits in model_rate_limits.items() - } + self.actual_rpm_limits: dict[str, int] = {model: limits["rpm"] for model, limits in model_rate_limits.items()} + self.actual_tpm_limits: dict[str, int] = {model: limits["tpm"] for model, limits in model_rate_limits.items()} # Enforced limits per model (with safety margin applied) self.enforced_rpm_limits: dict[str, int] = { - model: int(limit_safety_ratio * limits["rpm"]) - for model, limits in model_rate_limits.items() + model: int(limit_safety_ratio * limits["rpm"]) for model, limits in model_rate_limits.items() } self.enforced_tpm_limits: dict[str, int] = { - model: int(limit_safety_ratio * limits["tpm"]) - for model, limits in model_rate_limits.items() + model: int(limit_safety_ratio * limits["tpm"]) for model, limits in model_rate_limits.items() } # Request tracking @@ -337,8 +331,7 @@ async def get_current_usage(self): # Calculate total historical tokens for this model total_tokens = sum( - r.actual_tokens if r.actual_tokens is not None else r.projected_tokens - for r in model_all + r.actual_tokens if r.actual_tokens is not None else r.projected_tokens for r in model_all ) per_model_stats[model_name] = { @@ -395,9 +388,7 @@ async def get_current_usage(self): "average_requests_per_minute": (total_all_requests / session_duration * 60) if session_duration > 0 else 0, - "average_tokens_per_minute": (total_all_tokens / session_duration * 60) - if session_duration > 0 - else 0, + "average_tokens_per_minute": (total_all_tokens / session_duration * 60) if session_duration > 0 else 0, # Configuration "limit_safety_ratio": self.limit_safety_ratio, "minimum_wait_time": self.minimum_wait_time, diff --git a/rapidfireai/evals/utils/serialize.py b/rapidfireai/evals/utils/serialize.py index 6d1f2fa6..6200ffa9 100644 --- a/rapidfireai/evals/utils/serialize.py +++ b/rapidfireai/evals/utils/serialize.py @@ -4,7 +4,7 @@ import dill from rapidfireai.automl.model_config import RFvLLMModelConfig, RFOpenAIAPIModelConfig -from rapidfireai.evals.utils.constants import SEARCH_TYPE_KEYS +from rapidfireai.utils.constants import SEARCH_TYPE_KEYS def extract_pipeline_display_metadata(pipeline_config: dict[str, Any]) -> dict[str, Any]: @@ -173,9 +173,7 @@ def extract_pipeline_config_json(pipeline_config: dict[str, Any]) -> dict[str, A # Extract online_strategy_kwargs if present if "online_strategy_kwargs" in pipeline_config: - json_config["online_strategy_kwargs"] = pipeline_config[ - "online_strategy_kwargs" - ] + json_config["online_strategy_kwargs"] = pipeline_config["online_strategy_kwargs"] # Extract pipeline type and model-specific params if "pipeline" in pipeline_config: @@ -237,10 +235,7 @@ def extract_rag_params(rag_spec): json_config["pipeline_type"] = "openai" # Extract client_config (dict) - filter out sensitive keys - if ( - hasattr(pipeline, "client_config") - and pipeline.client_config is not None - ): + if hasattr(pipeline, "client_config") and pipeline.client_config is not None: sensitive_keys = {"api_key", "secret", "token", "password", "key"} json_config["client_config"] = { k: v for k, v in pipeline.client_config.items() @@ -252,10 +247,7 @@ def extract_rag_params(rag_spec): json_config["model_config"] = pipeline.model_config # Extract sampling_params using sampling_params_to_dict (extracts from model_config) - if ( - hasattr(pipeline, "sampling_params") - and pipeline.sampling_params is not None - ): + if hasattr(pipeline, "sampling_params") and pipeline.sampling_params is not None: json_config["sampling_params"] = pipeline.sampling_params_to_dict() # Extract rate limiting params @@ -263,10 +255,7 @@ def extract_rag_params(rag_spec): json_config["rpm_limit"] = pipeline.rpm_limit if hasattr(pipeline, "tpm_limit") and pipeline.tpm_limit is not None: json_config["tpm_limit"] = pipeline.tpm_limit - if ( - hasattr(pipeline, "max_completion_tokens") - and pipeline.max_completion_tokens is not None - ): + if hasattr(pipeline, "max_completion_tokens") and pipeline.max_completion_tokens is not None: json_config["max_completion_tokens"] = pipeline.max_completion_tokens # Extract RAG params if present @@ -281,4 +270,4 @@ def extract_rag_params(rag_spec): except (TypeError, ValueError) as e: raise ValueError(f"Failed to serialize pipeline config to JSON: {e}") from e - return json_config \ No newline at end of file + return json_config diff --git a/rapidfireai/experiment.py b/rapidfireai/experiment.py index 4243c583..5bc19291 100644 --- a/rapidfireai/experiment.py +++ b/rapidfireai/experiment.py @@ -1,31 +1,43 @@ """ -This module contains the unified Experiment class for both fit and evals modes. +RapidFire AI Experiment Module + +The Experiment class is the main entry point for running training (fit) or +inference (evals) experiments. """ import logging -import multiprocessing as mp import os import time import traceback from collections.abc import Callable -from typing import Any from pathlib import Path -from rapidfireai.utils.ping import ping_server +from typing import Any import pandas as pd + from rapidfireai.utils.constants import ( - ColabConfig, - RayConfig, - RF_EXPERIMENT_PATH, - RF_LOG_FILENAME, - RF_TRAINING_LOG_FILENAME, + ColabConfig, + DispatcherConfig, + ExperimentStatus, + RayConfig, + RF_EXPERIMENT_PATH, + RF_LOG_FILENAME, RF_LOG_PATH, - RF_MLFLOW_ENABLED + RF_MLFLOW_ENABLED, + RF_TRAINING_LOG_FILENAME, + get_dispatcher_url, ) +from rapidfireai.platform.ping import ping_server class Experiment: - """Unified Experiment class for both fit and evals modes.""" + """ + Main experiment class for RapidFire AI. + + Supports two modes: + - fit: Training experiments with hyperparameter search + - evals: Inference experiments with pipeline management + """ def __init__( self, @@ -59,21 +71,89 @@ def __init__( else: self._init_evals_mode() + def _init_ray(self) -> None: + """ + Connect to a running Ray cluster started by ``rapidfireai start``. + + If Ray is already initialized in this process (e.g., a previous experiment), + reuses the existing connection. On Colab, starts a local cluster automatically + since there is no terminal to run the CLI. + + Raises: + ConnectionError: If no Ray cluster is running (non-Colab only). + """ + import ray + + self._ray = ray + + os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") + os.environ.setdefault("VLLM_LOGGING_LEVEL", "ERROR") + os.environ.setdefault("RAY_LOG_TO_STDERR", "0") + os.environ.setdefault("MLFLOW_SUPPRESS_PRINTING_URL_TO_STDOUT", "true") + os.environ["RAY_DISABLE_IMPORT_WARNING"] = "1" + os.environ["RAY_DEDUP_LOGS"] = "0" + os.environ["RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO"] = "0" + + if ray.is_initialized(): + return + + ray_runtime_env = { + "env_vars": { + "CUDA_LAUNCH_BLOCKING": "0", + "CUDA_MODULE_LOADING": "LAZY", + "TF_CPP_MIN_LOG_LEVEL": "3", + "PYTORCH_CUDA_ALLOC_CONF": "max_split_size_mb:512", + "NCCL_NET": "Socket", + "NCCL_IB_DISABLE": "1", + } + } + + try: + ray.init( + address="auto", + ignore_reinit_error=True, + runtime_env=ray_runtime_env, + ) + except ConnectionError: + if not ColabConfig.ON_COLAB: + raise ConnectionError( + "No running RapidFire cluster found.\n" + "Please run `rapidfireai start` in a terminal before creating an experiment." + ) + ray.init( + logging_level=logging.ERROR, + log_to_driver=False, + configure_logging=True, + include_dashboard=True, + dashboard_host=RayConfig.HOST, + dashboard_port=RayConfig.PORT, + _metrics_export_port=None, + runtime_env=ray_runtime_env, + ) + + if ColabConfig.ON_COLAB: + try: + from google.colab.output import eval_js + + proxy_url = eval_js(f"google.colab.kernel.proxyPort({RayConfig.PORT})") + print(f"๐ŸŒ Google Colab detected. Ray dashboard URL: {proxy_url}") + except Exception as e: + if hasattr(self, "logger"): + self.logger.warning(f"Colab detected but failed to get proxy URL: {e}") + def _init_fit_mode(self) -> None: """Initialize fit-specific components.""" - # Import fit-specific modules - from rapidfireai.fit.db.rf_db import RfDb - from rapidfireai.fit.utils.exceptions import ExperimentException - from rapidfireai.fit.utils.experiment_utils import ExperimentUtils - from rapidfireai.fit.utils.logging import RFLogger + # Import unified modules + from rapidfireai.db.rf_db import RfDb + from rapidfireai.utils.exceptions import ExperimentException + from rapidfireai.utils.experiment_utils import ExperimentUtils + from rapidfireai.utils.logging import RFLogger from rapidfireai.version import __version__ # Store exception class for use in methods self._ExperimentException = ExperimentException # Initialize fit-specific attributes - self.log_server_process: mp.Process | None = None - self.worker_processes: list[mp.Process] = [] self._training_thread: Any = None # Track background training thread (Colab only) # Create db tables @@ -99,54 +179,39 @@ def _init_fit_mode(self) -> None: # Create logger try: - self.logger = RFLogger().create_logger("experiment") + self.logging_manager = RFLogger( + experiment_name=self.experiment_name, + experiment_path=self.experiment_path, + ) + self.logger = self.logging_manager.get_logger("Experiment") for msg in log_messages: self.logger.info(msg) - # Log the version of rapidfireai that is running self.logger.info(f"Running RapidFire AI version {__version__}") except Exception as e: raise ExperimentException(f"Error creating logger: {e}, traceback: {traceback.format_exc()}") from e - # Setup signal handlers for graceful shutdown + # Initialize or connect to shared Ray cluster try: - self.experiment_utils.setup_signal_handlers(self.worker_processes) + self._init_ray() + self.logger.info("Ray initialized successfully") 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 + raise ExperimentException(f"Error initializing Ray: {e}, traceback: {traceback.format_exc()}") from e def _init_evals_mode(self) -> None: """Initialize evals-specific components.""" - # Import evals-specific modules - import ray + # Import unified modules + from rapidfireai.db.rf_db import RfDb + from rapidfireai.dispatcher import start_dispatcher_thread + from rapidfireai.utils.experiment_utils import ExperimentUtils + from rapidfireai.utils.logging import RFLogger + from rapidfireai.platform.colab import get_colab_auth_token + from rapidfireai.utils.constants import DispatcherConfig + from rapidfireai.metrics import RFMetricLogger - from rapidfireai.evals.db import RFDatabase - from rapidfireai.evals.dispatcher import start_dispatcher_thread + # Import evals-specific modules from rapidfireai.evals.scheduling.controller import Controller - from rapidfireai.utils.colab import get_colab_auth_token - from rapidfireai.utils.constants import DispatcherConfig - from rapidfireai.evals.utils.constants import get_dispatcher_url - from rapidfireai.evals.utils.experiment_utils import ExperimentUtils - from rapidfireai.evals.utils.logger import RFLogger - from rapidfireai.utils.metric_rfmetric_manager import RFMetricLogger from rapidfireai.evals.utils.notebook_ui import NotebookUI - # Store ray reference for later use - self._ray = ray - - # Disable tokenizers parallelism warning when using with Ray/multiprocessing - os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") - # Suppress verbose third-party library logging - os.environ.setdefault("VLLM_LOGGING_LEVEL", "ERROR") - os.environ.setdefault("RAY_LOG_TO_STDERR", "0") - os.environ.setdefault("MLFLOW_SUPPRESS_PRINTING_URL_TO_STDOUT", "true") - # Disable Ray and other verbose logging - os.environ["RAY_DISABLE_IMPORT_WARNING"] = "1" - os.environ["RAY_DEDUP_LOGS"] = "0" - os.environ["RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO"] = "0" - # Initialize experiment utils self.experiment_utils = ExperimentUtils() @@ -164,40 +229,11 @@ def _init_evals_mode(self) -> None: for msg in log_messages: self.logger.info(msg) - # Initialize Ray with runtime environment for CUDA initialization - # This fixes AWS-specific CUDA/cuBLAS initialization issues - ray.init( - logging_level=logging.ERROR, - log_to_driver=False, - configure_logging=True, - ignore_reinit_error=True, - include_dashboard=True, - dashboard_host=RayConfig.HOST, - dashboard_port=RayConfig.PORT, - # Disable metrics export to prevent "Failed to establish connection" errors - _metrics_export_port=None, - runtime_env={ - "env_vars": { - # Force CUDA to initialize properly in Ray actors (AWS fix) - "CUDA_LAUNCH_BLOCKING": "0", - "CUDA_MODULE_LOADING": "LAZY", - "TF_CPP_MIN_LOG_LEVEL": "3", - "PYTORCH_CUDA_ALLOC_CONF": "max_split_size_mb:512", - } - }, - ) - if ColabConfig.ON_COLAB: - try: - from google.colab.output import eval_js - - # Get the Colab proxy URL for the dispatcher port - proxy_url = eval_js(f"google.colab.kernel.proxyPort({RayConfig.PORT})") - print(f"๐ŸŒ Google Colab detected. Ray dashboard URL: {proxy_url}") - except Exception as e: - print(f"โš ๏ธ Colab detected but failed to get proxy URL: {e}") + # Initialize or connect to shared Ray cluster + self._init_ray() # Create database reference - self.db = RFDatabase() + self.db = RfDb() try: metric_loggers = RFMetricLogger.get_default_metric_loggers(experiment_name=self.experiment_name) @@ -212,8 +248,6 @@ def _init_evals_mode(self) -> None: self.logger.warning(f"Failed to initialize MetricLogger: {e}. MetricLogger logging will be disabled.") self.metric_loggers = None - - # Initialize the controller self.controller = Controller( experiment_name=self.experiment_name, @@ -223,12 +257,17 @@ def _init_evals_mode(self) -> None: # Start dispatcher in background thread for interactive control if ping_server(DispatcherConfig.HOST, DispatcherConfig.PORT, 2): - self.logger.info(f"Using existing dispatcher/api server at {DispatcherConfig.HOST}:{DispatcherConfig.PORT}.") + self.logger.info( + f"Using existing dispatcher/api server at {DispatcherConfig.HOST}:{DispatcherConfig.PORT}." + ) self.dispatcher_thread = None - else: - self.logger.info(f"Starting new dispatcher/api server at {DispatcherConfig.HOST}:{DispatcherConfig.PORT}.") - self.dispatcher_thread = start_dispatcher_thread(host=DispatcherConfig.HOST, port=DispatcherConfig.PORT, logger=self.logger) + self.logger.info( + f"Starting new dispatcher/api server at {DispatcherConfig.HOST}:{DispatcherConfig.PORT}." + ) + self.dispatcher_thread = start_dispatcher_thread( + host=DispatcherConfig.HOST, port=DispatcherConfig.PORT, logger=self.logger + ) # Initialize notebook UI controller with auth token for Colab self.notebook_ui = NotebookUI(dispatcher_url=get_dispatcher_url(), auth_token=get_colab_auth_token()) @@ -262,14 +301,30 @@ def run_fit( raise ValueError("run_fit() is only available in 'fit' mode") from rapidfireai.fit.backend.controller import Controller - - ExperimentException = self._ExperimentException + from rapidfireai.utils.exceptions import ExperimentException # Check if training is already running if self._training_thread is not None and self._training_thread.is_alive(): print("โš ๏ธ Training is already running in background. Please wait for it to complete.") return + # Auto-detect resources from Ray cluster (similar to evals) + available_gpus = self._ray.cluster_resources().get("GPU", 0) + available_cpus = self._ray.cluster_resources().get("CPU", 0) + + # For fit mode: 1 GPU per worker (training requires full GPU) + num_workers = int(available_gpus) if available_gpus > 0 else 1 + gpus_per_worker = 1 if available_gpus > 0 else 0 + cpus_per_worker = max(1, int(available_cpus / num_workers)) if num_workers > 0 else 1 + + if available_gpus == 0: + self.logger.warning("No GPUs detected. Training will run on CPU (not recommended).") + else: + self.logger.info( + f"Detected {int(available_gpus)} GPU(s), {int(available_cpus)} CPU(s). " + f"Creating {num_workers} worker(s) with {gpus_per_worker} GPU(s) each." + ) + if ColabConfig.ON_COLAB: # Run Controller in background thread to keep kernel responsive import sys @@ -285,13 +340,19 @@ def _run_controller_background(): sys.stdout = StringIO() try: - controller = Controller(self.experiment_id, self.experiment_name) + controller = Controller( + self.experiment_id, + self.experiment_name, + num_workers=num_workers, + gpus_per_worker=gpus_per_worker, + cpus_per_worker=cpus_per_worker, + ) controller.run_fit(param_config, create_model_fn, train_dataset, eval_dataset, num_chunks, seed, num_gpus, monte_carlo_simulations) except Exception as e: # Restore stdout for error logging sys.stdout = old_stdout if hasattr(self, "logger"): - self.logger.opt(exception=True).error(f"Error in background training: {e}") + self.logger.exception(f"Error in background training: {e}") display(HTML(f'

โŒ Error in background training: {e}

')) finally: # Restore stdout @@ -324,11 +385,17 @@ def _run_controller_background(): else: # Original blocking behavior for non-Colab environments try: - controller = Controller(self.experiment_id, self.experiment_name) + controller = Controller( + self.experiment_id, + self.experiment_name, + num_workers=num_workers, + gpus_per_worker=gpus_per_worker, + cpus_per_worker=cpus_per_worker, + ) controller.run_fit(param_config, create_model_fn, train_dataset, eval_dataset, num_chunks, seed, num_gpus, monte_carlo_simulations) except Exception as e: if hasattr(self, "logger"): - self.logger.opt(exception=True).error(f"Error running fit: {e}") + self.logger.exception(f"Error running fit: {e}") raise ExperimentException(f"Error running fit: {e}, traceback: {traceback.format_exc()}") from e def run_evals( @@ -363,16 +430,14 @@ def run_evals( if self.mode != "evals": raise ValueError("run_evals() is only available in 'evals' mode") - from rapidfireai.evals.utils.constants import ExperimentStatus - # Auto-detect resources if not provided available_gpus = self._ray.cluster_resources().get("GPU", 0) available_cpus = self._ray.cluster_resources().get("CPU", 0) if gpus_per_actor is None: - gpus_per_actor = available_gpus if available_gpus > 1 else available_gpus/2 + gpus_per_actor = available_gpus if available_gpus > 1 else available_gpus / 2 if cpus_per_actor is None: - cpus_per_actor = available_cpus if available_cpus > 2 else available_cpus/2 + cpus_per_actor = available_cpus if available_cpus > 2 else available_cpus / 2 if num_actors is None: # Default to number of GPUs, or 1 if no GPUs available num_actors = int(gpus_per_actor) if gpus_per_actor > 0 else 1 @@ -426,8 +491,6 @@ def run_evals( self.db.set_experiment_error(self.experiment_id, str(e)) raise - - return results def get_results(self) -> pd.DataFrame: @@ -443,7 +506,7 @@ def get_results(self) -> pd.DataFrame: if self.mode != "fit": raise ValueError("get_results() is only available in 'fit' mode") - ExperimentException = self._ExperimentException + from rapidfireai.utils.exceptions import ExperimentException try: runs_info_df = self.experiment_utils.get_runs_info() @@ -458,7 +521,7 @@ def get_results(self) -> pd.DataFrame: return pd.DataFrame(columns=["run_id", "step"]) # Lazy import - only import when we actually have metric runs to fetch - from rapidfireai.utils.metric_rfmetric_manager import RFMetricLogger + from rapidfireai.metrics import RFMetricLogger try: metric_loggers = RFMetricLogger.get_default_metric_loggers(experiment_name=self.experiment_name) self.metric_loggers = RFMetricLogger(metric_loggers, logger=self.logger) @@ -494,7 +557,7 @@ def get_results(self) -> pd.DataFrame: except Exception as e: if hasattr(self, "logger"): - self.logger.opt(exception=True).error(f"Error getting results: {e}") + self.logger.exception(f"Error getting results: {e}") raise ExperimentException(f"Error getting results: {e}, traceback: {traceback.format_exc()}") from e def get_runs_info(self) -> pd.DataFrame: @@ -510,13 +573,13 @@ def get_runs_info(self) -> pd.DataFrame: if self.mode != "fit": raise ValueError("get_runs_info() is only available in 'fit' mode") - ExperimentException = self._ExperimentException + from rapidfireai.utils.exceptions import ExperimentException 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}") + self.logger.exception(f"Error getting run info: {e}") raise ExperimentException(f"Error getting run info: {e}, traceback: {traceback.format_exc()}") from e def cancel_current(self) -> None: @@ -525,53 +588,48 @@ def cancel_current(self) -> None: Works in both fit and evals modes. """ - if self.mode == "fit": - ExperimentException = self._ExperimentException - 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 - else: - # Eval mode + from rapidfireai.utils.exceptions import ExperimentException + + try: self.experiment_utils.cancel_current(internal=False) + except Exception as e: + if hasattr(self, "logger"): + self.logger.exception(f"Error canceling current task: {e}") + raise ExperimentException( + f"Error canceling current task: {e}, traceback: {traceback.format_exc()}" + ) from e def end(self) -> None: """ End the experiment and clean up resources. Works in both fit and evals modes with mode-specific cleanup. + Kills any Ray actors owned by this experiment to free GPU/CPU resources. + Ray cluster is intentionally kept alive so subsequent experiments + (fit or evals) can reuse it without port conflicts or restart delays. """ - if self.mode == "fit": - ExperimentException = self._ExperimentException + from rapidfireai.utils.exceptions import ExperimentException - 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 - else: - # Eval mode - # Use experiment_utils to end the experiment properly + try: self.experiment_utils.end_experiment(internal=False) + except Exception as e: + if hasattr(self, "logger"): + self.logger.exception(f"Error ending experiment: {e}") + raise ExperimentException(f"Error ending experiment: {e}, traceback: {traceback.format_exc()}") from e + + # Kill fit workers to free GPU resources for subsequent experiments. + # Workers are named fit_worker_0, fit_worker_1, etc. by create_worker_actors(). + # This is a no-op if no workers exist (e.g., evals mode or already cleaned up). + if self._ray.is_initialized(): + for i in range(16): + try: + worker = self._ray.get_actor(f"fit_worker_{i}") + self._ray.kill(worker) + except ValueError: + continue - # Clean shutdown Ray - self._ray.shutdown() - self.logger.info("All actors shut down") + self.logger.info("Experiment ended. Ray cluster remains active for subsequent experiments.") + if self.mode == "evals": self.logger.info("Dispatcher will automatically shut down (daemon thread)") def get_log_file_path(self, log_type: str | None = None) -> Path: diff --git a/rapidfireai/fit/backend/controller.py b/rapidfireai/fit/backend/controller.py index 26df2669..4ece1d02 100644 --- a/rapidfireai/fit/backend/controller.py +++ b/rapidfireai/fit/backend/controller.py @@ -6,79 +6,92 @@ import time from collections.abc import Callable from logging import Logger -from pathlib import Path from pprint import pformat from typing import Any -import torch +import ray from torch.utils.data import Dataset -from rapidfireai.automl import AutoMLAlgorithm -from rapidfireai.utils.os_utils import mkdir_p +from rapidfireai.automl import AutoMLAlgorithm, get_flattened_config_leaf, get_runs +from rapidfireai.db import RfDb from rapidfireai.fit.backend.chunks import DatasetChunks from rapidfireai.fit.backend.scheduler import Scheduler -from rapidfireai.fit.db.rf_db import RfDb -from rapidfireai.automl import get_flattened_config_leaf, get_runs -from rapidfireai.fit.utils.constants import ( +from rapidfireai.fit.backend.worker_actor import create_worker_actors +from rapidfireai.fit.utils.datapaths import DataPath +from rapidfireai.fit.utils.shm_manager import SharedMemoryManager +from rapidfireai.utils.constants import ( ControllerTask, ExperimentTask, + ICOperation, + ICStatus, RunEndedBy, RunSource, RunStatus, TaskStatus, WorkerTask, ) -from rapidfireai.fit.utils.datapaths import DataPath from rapidfireai.utils.distributed_utils import find_free_port -from rapidfireai.fit.utils.exceptions import ControllerException, NoGPUsFoundException -from rapidfireai.fit.utils.logging import RFLogger -from rapidfireai.utils.metric_rfmetric_manager import RFMetricLogger -from rapidfireai.fit.utils.serialize import encode_payload -from rapidfireai.fit.utils.shm_manager import SharedMemoryManager -from rapidfireai.fit.utils.worker_manager import WorkerManager +from rapidfireai.utils.exceptions import ControllerException, NoGPUsFoundException +from rapidfireai.utils.logging import RFLogger +from rapidfireai.metrics import RFMetricLogger +from rapidfireai.utils.os_utils import mkdir_p +from rapidfireai.utils.serialize import encode_payload class Controller: """This module contains the ML Controller class which is responsible for orchestrating the RapidFire lifecycle.""" - def __init__(self, experiment_id: int, experiment_name: str) -> None: - """Initialize the controller.""" - import torch.multiprocessing as mp - - with contextlib.suppress(RuntimeError): - mp.set_start_method("spawn", force=True) - + def __init__( + self, + experiment_id: int, + experiment_name: str, + num_workers: int = 1, + gpus_per_worker: int = 1, + cpus_per_worker: int = 1, + ) -> None: + """ + Initialize the controller. + + Args: + experiment_id: ID of the experiment + experiment_name: Name of the experiment + num_workers: Number of worker actors to create (typically = num GPUs) + gpus_per_worker: Number of GPUs per worker (typically 1) + cpus_per_worker: Number of CPUs per worker + """ self.experiment_id: int = experiment_id self.experiment_name: str = experiment_name + self.num_workers: int = num_workers + self.gpus_per_worker: int = gpus_per_worker + self.cpus_per_worker: int = cpus_per_worker - # create database object + # Create database object self.db: RfDb = RfDb() - # create controller logger - logging = RFLogger() - self.logger: Logger = logging.create_logger("controller") - self.user_logger: Logger = logging.create_logger("user") - self.ic_logger: Logger = logging.create_logger("interactive-control") + # Initialize data paths + DataPath.initialize(experiment_name, self.db.get_experiments_path(experiment_id)) + + # Create controller logger + logging = RFLogger(experiment_name=experiment_name) + self.logger: Logger = logging.get_logger("controller") + self.user_logger: Logger = logging.get_logger("user") + self.ic_logger: Logger = logging.get_logger("interactive-control") - # get number of GPUs - self.num_workers: int = torch.cuda.device_count() - if self.num_workers == 0: - raise NoGPUsFoundException("No GPUs found while initializing controller.") - self.logger.debug(f"Found {self.num_workers} workers/GPUs.") + # Validate GPU availability + if self.num_workers == 0 or (self.gpus_per_worker == 0 and self.num_workers > 0): + self.logger.warning("No GPUs available. Training will run on CPU (not recommended).") + else: + self.logger.debug(f"Configured for {self.num_workers} workers with {self.gpus_per_worker} GPU(s) each.") # set default required workers self.default_req_workers: int = 1 - # initialize shared manager and registry, create shared memory manager instance - self.shm_manager: SharedMemoryManager = SharedMemoryManager( - name="controller-shm" - ) - registry, process_lock = self.shm_manager.get_shm_objects() + # Initialize shared memory manager (used for model/checkpoint sharing between workers) + # SharedMemoryManager creates a RegistryActor internally for coordination + self.shm_manager: SharedMemoryManager = SharedMemoryManager(name="controller-shm") - # create worker manager - self.worker_manager: WorkerManager = WorkerManager( - self.num_workers, registry, process_lock - ) + # Worker actors (created in run_fit) + self.worker_actors: list = [] default_metric_loggers = RFMetricLogger.get_default_metric_loggers( experiment_name=self.experiment_name @@ -274,7 +287,7 @@ def _process_interactive_control( # process non-clone_modify tasks for run_id, run_state in run_states.items(): - if not run_state["task_id"]: + if not run_state["task"]: continue if run_state["status"] == RunStatus.STOPPED: @@ -285,9 +298,7 @@ def _process_interactive_control( status=RunStatus.STOPPED, ended_by=RunEndedBy.INTERACTIVE_CONTROL, ) - self.db.set_ic_ops_task_status( - run_state["task_id"], TaskStatus.COMPLETED - ) + self.db.update_ic_operation_status(run_state["ic_id"], ICStatus.COMPLETED) self.ic_logger.info(f"Stopping run {run_id} by Interactive Control") elif run_state["status"] == RunStatus.DELETED: # process deleted tasks @@ -295,6 +306,9 @@ def _process_interactive_control( # TODO: commented out to prevent clone of deleted runs issue (see Issue # 22) # self._clear_run_from_shm(run_id) + # cancel any scheduled worker tasks so workers don't pick them up + self.db.cancel_worker_tasks_for_run(run_id) + # delete run from MetricLogger metric_run_id = self.db.get_run(run_id)["metric_run_id"] self.metric_logger.delete_run(metric_run_id) @@ -304,9 +318,7 @@ def _process_interactive_control( status=RunStatus.DELETED, ended_by=RunEndedBy.INTERACTIVE_CONTROL, ) - self.db.set_ic_ops_task_status( - run_state["task_id"], TaskStatus.COMPLETED - ) + self.db.update_ic_operation_status(run_state["ic_id"], ICStatus.COMPLETED) self.ic_logger.info(f"Deleting run {run_id} by Interactive Control") elif run_state["status"] == RunStatus.ONGOING: # process ongoing tasks @@ -315,26 +327,19 @@ def _process_interactive_control( status=RunStatus.ONGOING, ended_by="", ) - self.db.set_ic_ops_task_status( - run_state["task_id"], TaskStatus.COMPLETED - ) + self.db.update_ic_operation_status(run_state["ic_id"], ICStatus.COMPLETED) self.ic_logger.info(f"Resuming run {run_id} by Interactive Control") elif run_state["status"] == RunStatus.COMPLETED: - # process completed tasks - self.logger.warning( - f"Run {run_id} is already completed. Skipping Interactive Control task." - ) - self.db.set_ic_ops_task_status(run_state["task_id"], TaskStatus.SKIPPED) + self.logger.warning(f"Run {run_id} is already completed. Skipping Interactive Control task.") + self.db.update_ic_operation_status(run_state["ic_id"], ICStatus.SKIPPED) else: raise ValueError(f"Unsupported run status {run_state['status']}") # process clone_modify tasks from the collected list for task in clone_modify_tasks: - parent_run_id, ic_op, config_leaf = ( - task["run_id"], - task["ic_op"], - task["config_leaf"], - ) + parent_run_id = task["target_id"] + operation = task["operation"] + config_leaf = task["config_data"] # add additional_kwargs to config_leaf if it exists in the parent run parent_run_details = self.db.get_run(parent_run_id) @@ -345,7 +350,7 @@ def _process_interactive_control( # create model for the new run try: - if ic_op == ControllerTask.IC_CLONE_MODIFY: + if operation == ICOperation.CLONE.value: clone_modify_info = { "cloned_from": parent_run_id, } @@ -357,7 +362,7 @@ def _process_interactive_control( num_chunks=num_chunks, clone_modify_info=clone_modify_info, ) - elif ic_op == ControllerTask.IC_CLONE_MODIFY_WARM: + elif operation == ICOperation.CLONE_WARM.value: # calculate clone chunk offset effective_batch_size = ( parent_run_details["config_leaf"]["training_args"].get( @@ -400,21 +405,17 @@ def _process_interactive_control( clone_modify_info, ) else: - raise ValueError(f"Unsupported IC operation {ic_op}") + raise ValueError(f"Unsupported IC operation {operation}") # mark task as completed - self.db.set_ic_ops_task_status(task["task_id"], TaskStatus.COMPLETED) + self.db.update_ic_operation_status(task["ic_id"], ICStatus.COMPLETED) self.ic_logger.info( - f"Cloned run {parent_run_id} by Interactive Control with {ic_op.value} into runs - {run_ids}" + f"Cloned run {parent_run_id} by Interactive Control with {operation} into runs - {run_ids}" ) except Exception as e: - self.db.set_ic_ops_task_status(task["task_id"], TaskStatus.FAILED) - self.ic_logger.error( - f"Error creating model for run {parent_run_id}: {e}" - ) - raise ControllerException( - f"Error creating model for run {parent_run_id}: {e}" - ) from e + self.db.update_ic_operation_status(task["ic_id"], ICStatus.FAILED) + self.ic_logger.error(f"Error creating model for run {parent_run_id}: {e}") + raise ControllerException(f"Error creating model for run {parent_run_id}: {e}") from e def _process_interm_ic_ops_states( self, @@ -422,22 +423,22 @@ def _process_interm_ic_ops_states( ) -> tuple[dict[str, Any], list[dict[str, Any]]]: """Process the interactive control.""" # get IC Ops scheduled tasks - ic_scheduled_tasks = self.db.get_scheduled_ic_ops_tasks() + ic_scheduled_tasks = self.db.get_pending_ic_operations(target_type="run") # track states for each task(run) and collect clone_modify tasks separately run_states = {} clone_modify_tasks = [] for task in ic_scheduled_tasks: - run_id = task["run_id"] + run_id = task["target_id"] # target_id is the run_id for run-targeted operations # skip if run is currently scheduled (we process IC ops only at chunk boundaries) if run_id in currently_scheduled_runs: # self.logger.debug(f"Skipping IC op for run {run_id} as it is currently scheduled") continue - is_clone_modify_task = task["ic_op"] in ( - ControllerTask.IC_CLONE_MODIFY, - ControllerTask.IC_CLONE_MODIFY_WARM, + is_clone_modify_task = task["operation"] in ( + ICOperation.CLONE.value, + ICOperation.CLONE_WARM.value, ) if is_clone_modify_task: @@ -452,77 +453,56 @@ def _process_interm_ic_ops_states( # track clone_modify tasks only for non-deleted runs if run_status != RunStatus.DELETED: clone_modify_tasks.append(task) - self.ic_logger.info(f"Added {task['ic_op']} task for run {run_id}.") + self.ic_logger.info(f"Added {task['operation']} task for run {run_id}.") else: - self.db.set_ic_ops_task_status(task["task_id"], TaskStatus.SKIPPED) - self.ic_logger.warning( - f"Skipping {task['ic_op']} task for deleted run {run_id}." - ) + self.db.update_ic_operation_status(task["ic_id"], ICStatus.SKIPPED) + self.ic_logger.warning(f"Skipping {task['operation']} task for deleted run {run_id}.") else: # Non clone_modify tasks if run_id not in run_states: run_states[run_id] = { - "task_id": None, + "ic_id": None, "task": None, "status": self.db.get_run(run_id)["status"], } # update run states based on existing status and task current_status = run_states[run_id]["status"] - if current_status == RunStatus.COMPLETED and task["ic_op"] in [ - ControllerTask.IC_RESUME, - ControllerTask.IC_STOP, + if current_status == RunStatus.COMPLETED and task["operation"] in [ + ICOperation.RESUME.value, + ICOperation.STOP.value, ]: # ignore RESUME/STOP tasks for completed runs - self.ic_logger.warning( - f"Ignoring RESUME/STOP task for run {run_id} as it is already completed" - ) - self.db.set_ic_ops_task_status(task["task_id"], TaskStatus.SKIPPED) - elif ( - current_status == RunStatus.FAILED - and task["ic_op"] != ControllerTask.IC_DELETE - ): - # ignore all tasks except DELETE for failed runs - self.ic_logger.warning( - f"Ignoring task {task['ic_op'].value} for failed run {run_id}" - ) - self.db.set_ic_ops_task_status(task["task_id"], TaskStatus.SKIPPED) + self.ic_logger.warning(f"Ignoring RESUME/STOP task for run {run_id} as it is already completed") + self.db.update_ic_operation_status(task["ic_id"], ICStatus.SKIPPED) + elif current_status == RunStatus.FAILED and task["operation"] != ICOperation.DELETE.value: + self.ic_logger.warning(f"Ignoring task {task['operation']} for failed run {run_id}") + self.db.update_ic_operation_status(task["ic_id"], ICStatus.SKIPPED) elif current_status == RunStatus.DELETED: - # ignore all tasks for deleted runs - self.ic_logger.warning( - f"Ignoring task {task['ic_op'].value} for deleted run {run_id}" - ) - self.db.set_ic_ops_task_status(task["task_id"], TaskStatus.SKIPPED) + self.ic_logger.warning(f"Ignoring task {task['operation']} for deleted run {run_id}") + self.db.update_ic_operation_status(task["ic_id"], ICStatus.SKIPPED) else: - # valid ic_op for this run - # mark prev task as completed - if run_states[run_id]["task_id"] is not None: - self.db.set_ic_ops_task_status( - run_states[run_id]["task_id"], TaskStatus.COMPLETED - ) + if run_states[run_id]["ic_id"] is not None: + self.db.update_ic_operation_status(run_states[run_id]["ic_id"], ICStatus.COMPLETED) # add new task to run states - if task["ic_op"] == ControllerTask.IC_STOP: + if task["operation"] == ICOperation.STOP.value: updated_status = RunStatus.STOPPED info_msg = f"Received STOP task for run {run_id}" - elif task["ic_op"] == ControllerTask.IC_DELETE: + elif task["operation"] == ICOperation.DELETE.value: updated_status = RunStatus.DELETED info_msg = f"Received DELETE task for run {run_id}" - elif task["ic_op"] == ControllerTask.IC_RESUME: + elif task["operation"] == ICOperation.RESUME.value: updated_status = RunStatus.ONGOING info_msg = f"Received RESUME task for run {run_id}" else: - self.db.set_ic_ops_task_status( - task["task_id"], TaskStatus.FAILED - ) - raise ValueError(f"Unsupported task {task['ic_op']}") + self.db.update_ic_operation_status(task["ic_id"], ICStatus.FAILED) + raise ValueError(f"Unsupported task {task['operation']}") run_states[run_id].update( { - "task_id": task["task_id"], - "task": task["ic_op"], - "status": ( - updated_status if updated_status else current_status - ), + "ic_id": task["ic_id"], + "task": task["operation"], + "status": (updated_status if updated_status else current_status), } ) self.ic_logger.info(info_msg) @@ -630,13 +610,22 @@ def run_fit( self.db.set_experiment_current_task(ExperimentTask.RUN_FIT) self.logger.debug(f"Set experiment task to {ExperimentTask.RUN_FIT.value}.") - # create workers + # Create Ray worker actors try: - self.worker_manager.create_workers() - print("Created workers") - self.logger.debug(f"Created {self.num_workers} workers.") + # Pass registry actor handle - workers use this for shared state coordination + self.worker_actors = create_worker_actors( + num_workers=self.num_workers, + gpus_per_worker=self.gpus_per_worker, + cpus_per_worker=self.cpus_per_worker, + registry_actor=self.shm_manager.get_registry_actor(), + ) + # Start all workers' serve_forever loops (fire-and-forget, runs in background) + for worker in self.worker_actors: + worker.serve_forever.remote() + self.logger.info(f"Created {self.num_workers} Ray worker actors") + self.logger.debug(f"Created {self.num_workers} Ray worker actors.") except Exception as e: - raise ControllerException(f"Error creating workers: {e}") from e + raise ControllerException(f"Error creating Ray worker actors: {e}") from e # create scheduler runs_info = [] @@ -666,7 +655,7 @@ def run_fit( while not all_done: # check for errors - exp_error = self.db.get_experiment_error() + exp_error = self.db.get_experiment_error(self.experiment_id) if exp_error: print(f"Error in experiment: {exp_error}") self.logger.error(f"Error in experiment: {exp_error}") @@ -904,7 +893,52 @@ def run_fit( self.logger.debug(f"Set experiment task to {ExperimentTask.IDLE.value}.") except Exception as e: + self._shutdown_workers() raise ControllerException(f"Error during run_fit: {e}") from e - # shutdown workers - self.worker_manager.shutdown() + # Shutdown workers gracefully + self._shutdown_workers() + + def _shutdown_workers(self, timeout: float = 60.0) -> None: + """ + Gracefully shutdown all worker actors. + + Uses hybrid approach: + 1. Request graceful shutdown (wait for chunk completion) + 2. ray.kill() as fallback after timeout + + Args: + timeout: Maximum seconds to wait for graceful shutdown + """ + if not self.worker_actors: + return + + self.logger.info("Shutting down worker actors...") + + # Request graceful shutdown for all workers + shutdown_futures = [] + for worker in self.worker_actors: + try: + shutdown_futures.append(worker.request_shutdown.remote()) + except Exception as e: + self.logger.warning(f"Error requesting shutdown: {e}") + + # Wait for graceful shutdown with timeout + try: + ray.get(shutdown_futures, timeout=timeout) + self.logger.info("All workers shut down gracefully") + except ray.exceptions.GetTimeoutError: + self.logger.warning(f"Timeout waiting for graceful shutdown after {timeout}s. Force killing workers...") + for worker in self.worker_actors: + try: + ray.kill(worker) + except Exception as e: + self.logger.debug(f"Error killing worker: {e}") + except Exception as e: + self.logger.warning(f"Error during graceful shutdown: {e}. Force killing workers...") + for worker in self.worker_actors: + with contextlib.suppress(Exception): + ray.kill(worker) + + self.worker_actors = [] + self.logger.info("Worker shutdown complete") diff --git a/rapidfireai/fit/utils/interactive_controller.py b/rapidfireai/fit/backend/interactive_controller.py similarity index 100% rename from rapidfireai/fit/utils/interactive_controller.py rename to rapidfireai/fit/backend/interactive_controller.py diff --git a/rapidfireai/fit/backend/worker.py b/rapidfireai/fit/backend/worker_actor.py similarity index 68% rename from rapidfireai/fit/backend/worker.py rename to rapidfireai/fit/backend/worker_actor.py index 2f3d9326..d90bdd01 100644 --- a/rapidfireai/fit/backend/worker.py +++ b/rapidfireai/fit/backend/worker_actor.py @@ -1,4 +1,9 @@ -"""This module contains the Worker class which is responsible for handling the worker operations.""" +""" +Ray Actor-based Worker for distributed training. + +This module provides a Ray remote class that replaces the multiprocessing-based Worker. +Each WorkerActor runs on a separate GPU and handles training tasks assigned by the Controller. +""" import gc import os @@ -8,16 +13,13 @@ from collections.abc import Callable from io import StringIO from logging import Logger -from multiprocessing import Process -from multiprocessing.managers import DictProxy -from multiprocessing.synchronize import Event as EventType -from multiprocessing.synchronize import Lock from typing import Any +import ray import torch +from rapidfireai.db import RfDb from rapidfireai.fit.backend.chunks import DatasetChunks -from rapidfireai.fit.db.rf_db import RfDb from rapidfireai.fit.ml.checkpoint_utils import ( flush_cuda_cache, purge_model_kv_caches, @@ -27,14 +29,15 @@ save_model_to_shared_memory, ) from rapidfireai.fit.ml.trainer import create_trainer_instance -from rapidfireai.fit.utils.constants import ( - USE_SHARED_MEMORY, +from rapidfireai.fit.utils.datapaths import DataPath +from rapidfireai.fit.utils.shm_manager import SharedMemoryManager, USE_SHARED_MEMORY +from rapidfireai.fit.utils.trainer_config import TrainerConfig +from rapidfireai.utils.constants import ( RunStatus, SHMObjectType, TaskStatus, WorkerTask, ) -from rapidfireai.fit.utils.datapaths import DataPath from rapidfireai.utils.distributed_utils import ( barrier, cleanup_distributed, @@ -42,17 +45,14 @@ setup_distributed_environment, find_free_port, ) -from rapidfireai.fit.utils.exceptions import WorkerException -from rapidfireai.fit.utils.logging import RFLogger, TrainingLogger -from rapidfireai.utils.metric_rfmetric_manager import RFMetricLogger -from rapidfireai.fit.utils.serialize import decode_db_payload -from rapidfireai.fit.utils.shm_manager import SharedMemoryManager -from rapidfireai.fit.utils.trainer_config import TrainerConfig +from rapidfireai.utils.exceptions import WorkerException +from rapidfireai.utils.logging import RFLogger, TrainingLogger +from rapidfireai.metrics import RFMetricLogger +from rapidfireai.utils.serialize import decode_db_payload def _scan_and_release_leaked_cuda_tensors(skip_quantized: bool = False): - """Zero all gc-tracked CUDA tensors unreachable by model iteration. - """ + """Zero all gc-tracked CUDA tensors unreachable by model iteration.""" gc.collect() gc.collect() @@ -68,7 +68,6 @@ def _scan_and_release_leaked_cuda_tensors(skip_quantized: bool = False): t = getattr(obj, attr_name, None) if torch.is_tensor(t): bnb_tensor_ids.add(id(t)) - # state2 is a nested QuantState for double quantization nested = getattr(obj, "state2", None) if isinstance(nested, QuantState): for attr_name in ("absmax", "code", "offset"): @@ -97,95 +96,85 @@ def __init__(self, original_stream, buffer): self.buffer = buffer def write(self, text): - """Write only to the buffer, not to the original stream (to suppress notebook output).""" self.buffer.write(text) - # Don't write to original_stream to prevent output in notebook def flush(self): - """Flush the buffer only.""" self.buffer.flush() - # Don't flush original_stream to prevent output in notebook def fileno(self): - """Return the original stream's file descriptor (required by vLLM).""" return self.original_stream.fileno() def isatty(self): - """Check if the original stream is a TTY.""" return self.original_stream.isatty() def getvalue(self): - """Get the captured buffer content.""" return self.buffer.getvalue() def __getattr__(self, name): - """Delegate other attributes to the original stream.""" return getattr(self.original_stream, name) -class Worker: - """Worker class that handles training and validation of runs""" +@ray.remote +class WorkerActor: + """ + Ray Actor-based Worker for handling training and validation. + + Each WorkerActor: + - Runs on its own GPU (resource allocation handled by Ray) + - Polls the database for scheduled tasks + - Executes training chunks when assigned + - Saves checkpoints to shared memory and disk + """ def __init__( self, worker_id: int, - model_registry: DictProxy, - process_lock: Lock, - shutdown_event: EventType, + registry_actor, ): - """Initialize the worker""" - self.process: Process - self.worker_id: int = worker_id - self.shutdown_event: EventType = shutdown_event + """ + Initialize the WorkerActor. - # Shared memory attributes (set by WorkerManager) - self.model_registry: DictProxy[int, Any] = model_registry - self.process_lock: Lock = process_lock + Args: + worker_id: Unique identifier for this worker + registry_actor: Ray actor handle for the shared RegistryActor + """ + self.worker_id: int = worker_id + self._shutdown_requested: bool = False - # Shared memory manager will be created using global objects + # Create shared memory manager using the registry actor + # The registry actor is shared across all workers for coordination self.shm_manager = SharedMemoryManager( name=f"worker-{worker_id}-shm", - registry=model_registry, - multiprocess_lock=process_lock, + registry_actor=registry_actor, ) - # create logger - self.logger: Logger = RFLogger().create_logger(f"worker_{worker_id}") - self.training_logger: Logger = TrainingLogger().create_logger( - f"worker_{worker_id}" - ) - self.logger.debug(f"Worker {self.worker_id} initialized with PID {os.getpid()}") - - # create database object self.db: RfDb = RfDb() - # get experiment name - self.experiment_name: str = self.db.get_running_experiment()["experiment_name"] + running_experiment = self.db.get_running_experiment() + self.experiment_name: str = running_experiment["experiment_name"] + self.experiment_id: int = running_experiment["experiment_id"] - # initialize data paths - DataPath.initialize( - self.experiment_name, self.db.get_experiments_path(self.experiment_name) - ) + self.logger: Logger = RFLogger(experiment_name=self.experiment_name).get_logger(f"worker_{worker_id}") + self.training_logger: Logger = TrainingLogger(experiment_name=self.experiment_name).get_logger(f"worker_{worker_id}") + self.logger.debug(f"WorkerActor {self.worker_id} initialized with PID {os.getpid()}") - # create metric logger - default_metric_loggers = RFMetricLogger.get_default_metric_loggers( - experiment_name=self.experiment_name - ) + DataPath.initialize(self.experiment_name, self.db.get_experiments_path(self.experiment_id)) + + default_metric_loggers = RFMetricLogger.get_default_metric_loggers(experiment_name=self.experiment_name) self.metric_logger = RFMetricLogger(default_metric_loggers, logger=self.logger) if self.metric_logger is None: - raise WorkerException( - "MetricLogger is not initialized. Please check the metric logger configuration." - ) + raise WorkerException("MetricLogger is not initialized. Please check the metric logger configuration.") self.metric_logger.get_experiment(self.experiment_name) - # load datasets - self.train_dataset, self.eval_dataset, self.num_chunks = self.load_datasets() + self.train_dataset, self.eval_dataset, self.num_chunks = self._load_datasets() self.len_train_dataset = len(self.train_dataset) - def load_datasets( + self.logger.info(f"WorkerActor {self.worker_id} ready on GPU {os.environ.get('CUDA_VISIBLE_DEVICES', 'N/A')}") + + def _load_datasets( self, ) -> tuple[torch.utils.data.Dataset | None, torch.utils.data.Dataset | None, int]: - """Load the train and eval datasets""" + """Load the train and eval datasets from disk.""" try: with open(DataPath.dataset_path(), "rb") as f: datasets = decode_db_payload(f.read()) @@ -201,22 +190,26 @@ def run_fit( multi_worker_details: dict[str, Any], create_model_fn: Callable, ) -> None: - """Run fit""" - self.logger.debug( - f"Received run_fit on worker for run {run_id} with chunk {chunk_id}" - ) + """ + Execute a training chunk for the specified run. + + Args: + run_id: ID of the training run + chunk_id: ID of the chunk to train on + multi_worker_details: FSDP distributed training configuration + create_model_fn: Function to create/load the model + """ + self.logger.debug(f"Received run_fit on worker for run {run_id} with chunk {chunk_id}") - # get run details run_details = self.db.get_run(run_id) config_leaf = run_details["config_leaf"] metric_run_id = run_details["metric_run_id"] - # check if FSDP is enabled use_fsdp = ( "training_args" in config_leaf and "fsdp_config" in config_leaf["training_args"] ) - # Clean up stale vLLM parallel state (critical for sequential GRPO training) + if config_leaf.get("trainer_type", "SFT") == "GRPO": try: from vllm.distributed import parallel_state @@ -225,35 +218,32 @@ def run_fit( try: parallel_state.destroy_model_parallel() except Exception: - pass # Ignore if not initialized + pass if hasattr(parallel_state, "destroy_distributed_environment"): try: parallel_state.destroy_distributed_environment() except Exception: - pass # Ignore if not initialized + pass except (ImportError, AttributeError): - pass # vLLM parallel state cleanup not available + pass - # Initialize distributed training if FSDP is enabled for this run if use_fsdp: try: os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" - # Get distributed configuration from run details master_addr = multi_worker_details["master_address"] master_port = multi_worker_details["master_port"] world_size = multi_worker_details["world_size"] rank = multi_worker_details["local_rank"] self.logger.debug( - f"Worker {self.worker_id} initializing distributed training for run {run_id} with master_addr {master_addr}, master_port {master_port}, world_size {world_size}, rank {rank}" - ) - os.environ["CUDA_VISIBLE_DEVICES"] = ",".join( - map(str, multi_worker_details["worker_ids"]) + f"Worker {self.worker_id} initializing distributed training for run {run_id} " + f"with master_addr {master_addr}, master_port {master_port}, world_size {world_size}, rank {rank}" ) setup_distributed_environment( rank=rank, world_size=world_size, master_addr=master_addr, master_port=master_port, + local_rank=0, ) self.logger.debug( f"Worker {self.worker_id} initialized distributed training for run {run_id}" @@ -269,31 +259,26 @@ def run_fit( master_port = find_free_port() os.environ["MASTER_PORT"] = str(master_port) - # get effective batch size - per_device_train_batch_size = config_leaf["training_args"].get( - "per_device_train_batch_size", 1 - ) - gradient_accumulation_steps = config_leaf["training_args"].get( - "gradient_accumulation_steps", 1 - ) + per_device_train_batch_size = config_leaf["training_args"].get("per_device_train_batch_size", 1) + gradient_accumulation_steps = config_leaf["training_args"].get("gradient_accumulation_steps", 1) effective_batch_size = ( per_device_train_batch_size * gradient_accumulation_steps * multi_worker_details.get("world_size", 1) ) - # fetch train dataset chunk train_dataset_chunker = DatasetChunks( self.len_train_dataset, self.num_chunks, batch_size=effective_batch_size, offset=run_details["chunk_offset"], ) - train_dataset_chunk = train_dataset_chunker.get_chunk( - self.train_dataset, chunk_id - ) + train_dataset_chunk = train_dataset_chunker.get_chunk(self.train_dataset, chunk_id) - # create worker config + # NOTE: local_rank here is the FSDP group rank (0..N-1), NOT the CUDA device + # index. Each Ray actor sees only its assigned GPU as cuda:0, so the device + # index is always 0. This rank is used for conditional logic (e.g., only rank 0 + # saves checkpoints) throughout checkpoint_utils.py and this file. trainer_config = TrainerConfig( worker_id=self.worker_id, run_id=run_id, @@ -311,27 +296,20 @@ def run_fit( cloned_from=run_details["cloned_from"], num_epochs_completed=run_details["num_epochs_completed"], ) - # add reward funcs to config_leaf if cloned from a GRPO run - if ( - trainer_config.cloned_from is not None - and trainer_config.config_leaf.get("trainer_type") == "GRPO" - ): + + if trainer_config.cloned_from is not None and trainer_config.config_leaf.get("trainer_type") == "GRPO": parent_run_details = self.db.get_run(trainer_config.cloned_from) - config_leaf["reward_funcs"] = parent_run_details["config_leaf"].get( - "reward_funcs" - ) + config_leaf["reward_funcs"] = parent_run_details["config_leaf"].get("reward_funcs") self.db.set_run_details(run_id, config_leaf=config_leaf) if use_fsdp and is_distributed_initialized(): barrier() - # Redirect stdout/stderr to /dev/null, capture Python-level output for logging stdout_buffer = StringIO() stderr_buffer = StringIO() original_stdout = sys.stdout original_stderr = sys.stderr - # Redirect to /dev/null to suppress notebook output devnull_fd = None try: with open(os.devnull, "w") as devnull: @@ -343,7 +321,6 @@ def run_fit( except Exception as e: self.logger.debug(f"Could not redirect stdout/stderr to /dev/null: {e}") - # TeeOutput captures to buffer; fileno() returns /dev/null fd for vLLM tee_stdout = TeeOutput(original_stdout, stdout_buffer) tee_stderr = TeeOutput(original_stderr, stderr_buffer) @@ -358,43 +335,37 @@ def run_fit( self.metric_logger, chunk_id, use_fsdp=use_fsdp, + db=self.db, ) is_quantized = bool( - config_leaf.get("model_kwargs", {}).get("quantization_config") or getattr(trainer_instance.model.config, "quantization_config", None) is not None + config_leaf.get("model_kwargs", {}).get("quantization_config") + or getattr(trainer_instance.model.config, "quantization_config", None) is not None ) - # if first time, save checkpoint to disk + completed_steps = self.db.get_completed_steps(run_id) if completed_steps == 0 and not USE_SHARED_MEMORY: save_checkpoint_to_disk( trainer_instance, trainer_config, first=True, use_fsdp=use_fsdp ) - # update base model name in db for run - trainer_config.config_leaf["model_name"] = ( - trainer_instance.model.config._name_or_path - ) + trainer_config.config_leaf["model_name"] = trainer_instance.model.config._name_or_path self.db.set_run_details(run_id, config_leaf=trainer_config.config_leaf) - self.logger.debug( - f"Beginning training for run {run_id} on chunk {chunk_id}" - ) + self.logger.debug(f"Beginning training for run {run_id} on chunk {chunk_id}") - # Synchronize all workers before training starts if use_fsdp and is_distributed_initialized(): barrier() start_time = time.time() trainer_instance.train() end_time = time.time() - # Synchronize all workers after training completes + if use_fsdp and is_distributed_initialized(): barrier() - # update completed steps new_completed_steps = completed_steps + trainer_instance.state.global_step self.db.set_completed_steps(run_id, new_completed_steps) - # update running average runtime in database for scheduler optimization if trainer_config.local_rank == 0: new_runtime_per_batch = ( end_time - start_time @@ -405,13 +376,10 @@ def run_fit( ) / new_completed_steps self.db.set_estimated_runtime(run_id, running_average_runtime) - # save checkpoints if USE_SHARED_MEMORY: if use_fsdp and is_distributed_initialized(): barrier() - self.logger.debug( - f"Worker {self.worker_id} passed barrier after training" - ) + self.logger.debug(f"Worker {self.worker_id} passed barrier after training") if use_fsdp and torch.cuda.is_available(): try: @@ -443,22 +411,13 @@ def run_fit( ) if is_run_finished: - # Skip shared memory for final chunk (flat_param zeroing would corrupt state) save_checkpoint_to_disk( - trainer_instance, - trainer_config, - last=True, - use_fsdp=use_fsdp, - ) - self.logger.debug( - f"Saved final checkpoint to disk for run {run_id} on chunk {chunk_id}" + trainer_instance, trainer_config, last=True, use_fsdp=use_fsdp, ) + self.logger.debug(f"Saved final checkpoint to disk for run {run_id} on chunk {chunk_id}") else: save_checkpoint_to_shared_memory( - trainer_instance, - trainer_config, - self.shm_manager, - use_fsdp=use_fsdp, + trainer_instance, trainer_config, self.shm_manager, use_fsdp=use_fsdp, ) if not config_leaf.get("peft_params") and not use_fsdp: save_model_to_shared_memory( @@ -471,40 +430,27 @@ def run_fit( use_fsdp=use_fsdp, ) self.logger.debug( - f"Saved checkpoint to shared memory for run {run_id} on chunk {chunk_id}", - f"and worker {self.worker_id}", + f"Saved checkpoint to shared memory for run {run_id} on chunk {chunk_id}" + f" and worker {self.worker_id}", ) if use_fsdp and is_distributed_initialized(): barrier() - # save checkpoint to disk based on save strategy - save_strategy = trainer_config.config_leaf.get("training_args", {}).get( - "save_strategy", "epoch" - ) + save_strategy = trainer_config.config_leaf.get("training_args", {}).get("save_strategy", "epoch") if save_strategy == "chunk": save_checkpoint_to_disk( - trainer_instance, - trainer_config, - completed_steps=new_completed_steps, - use_fsdp=use_fsdp, - ) - self.logger.debug( - f"Saved checkpoint to disk for run {run_id} on chunk {chunk_id}" + trainer_instance, trainer_config, + completed_steps=new_completed_steps, use_fsdp=use_fsdp, ) + self.logger.debug(f"Saved checkpoint to disk for run {run_id} on chunk {chunk_id}") else: - # save checkpoint to disk when not using shared memory save_checkpoint_to_disk( - trainer_instance, - trainer_config, - completed_steps=new_completed_steps, - use_fsdp=use_fsdp, - ) - self.logger.debug( - f"Saved checkpoint to disk for run {run_id} on chunk {chunk_id}" + trainer_instance, trainer_config, + completed_steps=new_completed_steps, use_fsdp=use_fsdp, ) + self.logger.debug(f"Saved checkpoint to disk for run {run_id} on chunk {chunk_id}") - # Save final checkpoint (non-shared-memory path only) if ( not USE_SHARED_MEMORY and chunk_id == self.num_chunks - 1 @@ -513,9 +459,7 @@ def run_fit( save_checkpoint_to_disk( trainer_instance, trainer_config, last=True, use_fsdp=use_fsdp ) - self.logger.debug( - f"Saved final checkpoint for run {run_id} on chunk {chunk_id}" - ) + self.logger.debug(f"Saved final checkpoint for run {run_id} on chunk {chunk_id}") if use_fsdp and is_distributed_initialized(): barrier() @@ -524,18 +468,14 @@ def run_fit( self.logger.error(f"Error during training for run {run_id}: {e}") raise finally: - # Cleanup must happen even when training fails to prevent memory leaks try: - # Clean up vLLM engine (critical for GRPO with sequential model training) if "trainer_instance" in locals() and hasattr(trainer_instance, "llm"): try: if hasattr(trainer_instance.llm, "shutdown"): trainer_instance.llm.shutdown() elif hasattr(trainer_instance.llm, "llm_engine"): if hasattr(trainer_instance.llm.llm_engine, "engine_core"): - if hasattr( - trainer_instance.llm.llm_engine.engine_core, "shutdown" - ): + if hasattr(trainer_instance.llm.llm_engine.engine_core, "shutdown"): trainer_instance.llm.llm_engine.engine_core.shutdown() try: from vllm.distributed import parallel_state @@ -555,7 +495,6 @@ def run_fit( except Exception: pass - # Systematic GPU memory cleanup if "trainer_instance" in locals(): _is_quantized_fsdp = ( use_fsdp and "is_quantized" in dir() and is_quantized @@ -623,9 +562,7 @@ def run_fit( if torch.cuda.is_available(): if use_fsdp and not _is_quantized_fsdp: _scan_and_release_leaked_cuda_tensors( - skip_quantized=( - "is_quantized" in dir() and is_quantized - ) + skip_quantized=("is_quantized" in dir() and is_quantized) ) gc.collect() torch.cuda.empty_cache() @@ -639,25 +576,27 @@ def run_fit( sys.stdout = original_stdout sys.stderr = original_stderr - # Write captured output to training logger if stdout_buffer.getvalue(): self.training_logger.info(stdout_buffer.getvalue()) if stderr_buffer.getvalue(): self.training_logger.error(stderr_buffer.getvalue()) def serve_forever(self) -> None: - """The main loop for the worker""" + """ + Main loop: poll the database for scheduled tasks and execute them. + This method runs continuously until shutdown is requested (via request_shutdown()) + or a critical error occurs. + """ prev_task_id: int | None = None - while not (self.shutdown_event and self.shutdown_event.is_set()): + + while not self._shutdown_requested: try: scheduled_task = self.db.get_worker_scheduled_task(self.worker_id) if not scheduled_task or scheduled_task["task_id"] == prev_task_id: - # no new tasks or same task as previous iteration time.sleep(1) continue - # get task details prev_task_id = scheduled_task["task_id"] task_type = scheduled_task["task_type"] run_id = scheduled_task["run_id"] @@ -667,20 +606,13 @@ def serve_forever(self) -> None: self.logger.debug(f"Received task {task_type} for run {run_id}") if task_type == WorkerTask.TRAIN_VAL: - self.db.set_worker_task_status( - self.worker_id, TaskStatus.IN_PROGRESS - ) + self.db.set_worker_task_status(self.worker_id, TaskStatus.IN_PROGRESS) - # run train and validation function try: - self.run_fit( - run_id, chunk_id, multi_worker_details, create_model_fn - ) - self.db.set_worker_task_status( - self.worker_id, TaskStatus.COMPLETED - ) + self.run_fit(run_id, chunk_id, multi_worker_details, create_model_fn) + self.db.set_worker_task_status(self.worker_id, TaskStatus.COMPLETED) except Exception as e: - self.logger.opt(exception=True).error( + self.logger.exception( f"Error while running run_fit for run {run_id} and chunk {chunk_id}: {e}" ) self.db.set_run_details( @@ -688,48 +620,93 @@ def serve_forever(self) -> None: status=RunStatus.FAILED, error=str(e) + traceback.format_exc(), ) - self.db.set_worker_task_status( - self.worker_id, TaskStatus.FAILED - ) + self.db.set_worker_task_status(self.worker_id, TaskStatus.FAILED) else: raise WorkerException(f"Invalid task type: {task_type}") if is_distributed_initialized(): cleanup_distributed() - self.logger.debug( - f"Worker {self.worker_id} distributed training cleaned up" - ) + self.logger.debug(f"Worker {self.worker_id} distributed training cleaned up") except Exception as e: - self.logger.opt(exception=True).error( - f"Worker {self.worker_id} error: {e}" - ) - self.db.set_experiment_error(str(e) + "\n" + traceback.format_exc()) + self.logger.exception(f"WorkerActor {self.worker_id} error: {e}") + self.db.set_experiment_error(self.experiment_id, str(e) + "\n" + traceback.format_exc()) break - self.shutdown() + self._cleanup() + + def request_shutdown(self) -> bool: + """ + Request graceful shutdown of this worker. + + The worker will complete its current chunk before shutting down. + + Returns: + True if shutdown was requested successfully + """ + self.logger.debug(f"WorkerActor {self.worker_id} shutdown requested") + self._shutdown_requested = True + return True + + def is_shutdown_requested(self) -> bool: + """Check if shutdown has been requested.""" + return self._shutdown_requested - def shutdown(self): - """Called by WorkerManager to gracefully shutdown this worker""" - self.logger.debug(f"Worker {self.worker_id} shutdown requested") - if self.shutdown_event: - self.shutdown_event.set() + def _cleanup(self) -> None: + """Clean up resources on shutdown.""" + self.logger.debug(f"WorkerActor {self.worker_id} cleaning up") - # Clean up distributed training if enabled if is_distributed_initialized(): try: cleanup_distributed() - self.logger.debug( - f"Worker {self.worker_id} distributed training cleaned up" - ) + self.logger.debug(f"WorkerActor {self.worker_id} distributed training cleaned up") except Exception as e: self.logger.debug(f"Error during distributed cleanup: {e}") - # Close database connection to prevent resource leaks try: if hasattr(self, "db"): self.db.close() except Exception as e: self.logger.debug(f"Error closing database connection: {e}") - def is_alive(self): - """Check if the worker process is alive""" - return self.process and self.process.is_alive() + self.logger.info(f"WorkerActor {self.worker_id} shutdown complete") + + def get_worker_id(self) -> int: + """Return this worker's ID.""" + return self.worker_id + + def ping(self) -> bool: + """Health check endpoint.""" + return True + + +def create_worker_actors( + num_workers: int, + gpus_per_worker: int, + cpus_per_worker: int, + registry_actor, +) -> list: + """ + Create WorkerActor instances with proper GPU allocation. + + Args: + num_workers: Number of workers to create + gpus_per_worker: Number of GPUs per worker (typically 1) + cpus_per_worker: Number of CPUs per worker + registry_actor: Ray actor handle for the shared RegistryActor + + Returns: + List of WorkerActor handles + """ + workers = [] + + for worker_id in range(num_workers): + worker = WorkerActor.options( + num_gpus=gpus_per_worker, + num_cpus=cpus_per_worker, + name=f"fit_worker_{worker_id}", + ).remote( + worker_id=worker_id, + registry_actor=registry_actor, + ) + workers.append(worker) + + return workers diff --git a/rapidfireai/fit/db/CLAUDE.md b/rapidfireai/fit/db/CLAUDE.md deleted file mode 100644 index f4551bb8..00000000 --- a/rapidfireai/fit/db/CLAUDE.md +++ /dev/null @@ -1,249 +0,0 @@ -# CLAUDE.md - Database - -This file provides guidance for working with the database layer of RapidFire AI. - -## Overview - -The db module provides the persistence layer for RapidFire using SQLite. It stores experiment metadata, run configurations, task scheduling state, and checkpoint locations. The design emphasizes async operations and clean separation between the database interface and domain logic. - -## Files - -### rf_db.py -**Purpose**: High-level database API with domain-specific operations - -**Key Responsibilities**: -- Implements all CRUD operations for experiments, runs, tasks, and IC Ops -- Handles serialization/deserialization of complex objects (using `encode_payload`/`decode_db_payload`) -- Manages experiment lifecycle (create, activate, reset, cleanup) -- Provides query methods for Controller, Worker, and Dispatcher -- Enforces business logic constraints (e.g., can't IC Ops on KILLED runs) - -**Key Methods - Experiments**: -- `create_experiment()`: Create new experiment entry -- `get_running_experiment()`: Get currently active experiment -- `set_experiment_status()`: Update experiment status -- `reset_all_tables()`: Truncate tables (cleanup) -- `reset_experiment_states()`: Mark in-progress tasks as failed (crash recovery) - -**Key Methods - Runs**: -- `create_run()`: Create run with config, status, source -- `get_run()`: Get run by ID -- `get_all_runs()`: Get all runs with metrics -- `get_runs_by_status()`: Filter runs by status(es) -- `set_run_status()`: Update run status -- `set_run_ended_by()`: Mark how run ended (completed/failed/killed) -- `update_run_metrics()`: Update training metrics - -**Key Methods - Tasks**: -- `create_worker_task()`: Create task for worker to execute -- `get_next_worker_task()`: Poll for next task (used by Worker) -- `set_worker_task_status()`: Update task status -- `get_controller_task()`: Get current controller task -- `set_controller_task()`: Update controller task - -**Key Methods - Interactive Control**: -- `request_clone_modify()`: Create IC Ops request for clone -- `request_stop()`: Request run stop -- `request_resume()`: Request run resume -- `request_delete()`: Request run deletion -- `get_ic_ops_request()`: Poll for IC Ops requests (used by Controller) -- `mark_ic_ops_completed()`: Mark IC Op as processed - -**Serialization**: -- Complex objects (configs, datasets, models) stored as BLOBs -- Uses `encode_payload()` from utils/serialize.py before storing -- Uses `decode_db_payload()` when reading back -- Handles torch tensors, datasets, and arbitrary Python objects via dill - -**Patterns**: -- All methods use `db.execute()` with parameterized queries (SQL injection safe) -- Commit=True by default for most operations -- Returns dicts or lists of dicts (not raw tuples) -- Raises DBException on errors with context - -### db_interface.py -**Purpose**: Low-level SQLite connection wrapper - -**Key Responsibilities**: -- Manages SQLite connection lifecycle -- Provides generic `execute()` method for queries -- Handles connection pooling and thread safety -- Converts query results to dicts - -**Key Methods**: -- `execute()`: Execute parameterized query, return results as list of dicts -- `close()`: Close database connection - -**Design Notes**: -- Uses sqlite3 row_factory for dict results -- Single connection per RfDb instance -- Thread-safe via SQLite's default settings - -### tables.sql -**Purpose**: Database schema definition - -**Tables**: - -**experiments**: -- `experiment_id` (PK): Unique experiment identifier -- `experiment_name`: User-provided name -- `status`: ExperimentStatus enum (NEW, RUNNING, COMPLETED, FAILED) -- `experiments_path`: Base path for experiment artifacts -- `created_at`, `updated_at`: Timestamps - -**runs**: -- `run_id` (PK): Unique run identifier -- `experiment_id` (FK): Parent experiment -- `run_name`: Generated name (e.g., "run_1") -- `metric_run_id`: Metric tracking ID -- `status`: RunStatus enum (NEW, ONGOING, COMPLETED, FAILED, STOPPED, KILLED) -- `source`: RunSource enum (USER, CLONE_MODIFY) -- `ended_by`: RunEndedBy enum (COMPLETED, FAILED, KILLED, STOPPED) -- `parent_run_id`: For cloned runs -- `warm_start`: Boolean flag for clone-modify -- `config_leaf`: Serialized run configuration (BLOB) -- `seed`: Random seed for reproducibility -- `num_chunks`: Number of data chunks -- `current_chunk`, `current_epoch`: Progress tracking -- `metrics`: JSON string of training metrics -- `created_at`, `updated_at`: Timestamps - -**worker_task**: -- `task_id` (PK): Unique task identifier -- `run_id` (FK): Run to execute -- `worker_id`: GPU worker assignment -- `chunk_id`: Data chunk to train on -- `epoch`: Current epoch number -- `status`: TaskStatus enum (SCHEDULED, IN_PROGRESS, COMPLETED, FAILED) -- `created_at`, `updated_at`: Timestamps - -**controller_progress**: -- Tracks controller state (single row table) -- `task`: ControllerTask enum -- `status`: TaskStatus enum - -**worker_progress**: -- Tracks per-worker state -- `worker_id` (PK): Worker identifier -- `task`: WorkerTask enum -- `status`: TaskStatus enum - -**interactive_control**: -- `ic_id` (PK): IC Ops request identifier -- `run_id` (FK): Target run -- `operation`: IC Ops type (CLONE_MODIFY, STOP, RESUME, DELETE) -- `status`: TaskStatus enum -- `config_leaf`: New config for clone-modify (BLOB) -- `warm_start`: Boolean for clone-modify -- `created_at`, `updated_at`: Timestamps - -## Key Concepts - -### Status Enums -Defined in `utils/constants.py`: -- **ExperimentStatus**: NEW, RUNNING, COMPLETED, FAILED -- **RunStatus**: NEW, ONGOING, COMPLETED, FAILED, STOPPED, KILLED -- **RunSource**: USER, CLONE_MODIFY -- **RunEndedBy**: COMPLETED, FAILED, KILLED, STOPPED -- **TaskStatus**: PENDING, SCHEDULED, IN_PROGRESS, COMPLETED, FAILED - -### Transaction Model -- Most operations are auto-commit (commit=True) -- No explicit transaction management (SQLite handles implicitly) -- Crash recovery via `reset_experiment_states()` on startup - -### Query Patterns -```python -# Parameterized query (safe) -query = "SELECT * FROM runs WHERE run_id = ?" -result = self.db.execute(query, (run_id,)) - -# With commit -query = "UPDATE runs SET status = ? WHERE run_id = ?" -self.db.execute(query, (new_status, run_id), commit=True) -``` - -## Common Operations - -### Creating a Run -```python -run_id = db.create_run( - experiment_id=1, - run_name="run_1", - metric_run_id="abc123", - config_leaf=encode_payload(config_dict), - source=RunSource.USER, - seed=42, - num_chunks=8 -) -``` - -### Polling for Tasks (Worker) -```python -task = db.get_next_worker_task(worker_id=0) -if task: - db.set_worker_task_status(task['task_id'], TaskStatus.IN_PROGRESS) - # ... do work ... - db.set_worker_task_status(task['task_id'], TaskStatus.COMPLETED) -``` - -### IC Ops Flow (Controller) -```python -# User triggers stop via UI โ†’ Dispatcher โ†’ Database -db.request_stop(run_id=5) - -# Controller polls and processes -ic_ops = db.get_ic_ops_request() -for op in ic_ops: - if op['operation'] == 'STOP': - # ... handle stop ... - db.mark_ic_ops_completed(op['ic_id']) -``` - -### Cleanup Between Experiments -```python -db.reset_all_tables(experiments_table=False) # Keep experiments table -db.set_experiment_status(exp_id, ExperimentStatus.COMPLETED) -``` - -## Testing Database Changes - -1. Modify `tables.sql` if adding/changing tables -2. Delete existing `rapidfire.db` file to force recreation -3. Run `db.create_tables()` to apply schema -4. Test with `pytest` or manual verification -5. Ensure backward compatibility with existing experiments - -## Performance Considerations - -- SQLite write contention: Workers only write task status updates -- Most writes are from Controller (runs, IC Ops, metrics) -- No indexes beyond PRIMARY KEYs (small data volume) -- BLOB storage for configs is fine (not queried, only retrieved by PK) - -## Common Patterns - -### Adding New Table -1. Add CREATE TABLE to `tables.sql` -2. Add CRUD methods to `rf_db.py` -3. Add any enums to `utils/constants.py` -4. Update `reset_all_tables()` if needed for cleanup -5. Test with fresh database - -### Debugging Database Issues -```python -# Check database file location -import os -print(os.path.abspath("rapidfire.db")) - -# Inspect directly with sqlite3 CLI -# sqlite3 rapidfire.db -# .schema -# SELECT * FROM runs; -# SELECT * FROM worker_task WHERE status = 'IN_PROGRESS'; -``` - -### Migration Strategy -- Currently no migrations (schema assumed stable) -- Breaking changes require users to reset experiments -- Future: Add migration system (e.g., Alembic) if needed diff --git a/rapidfireai/fit/db/__init__.py b/rapidfireai/fit/db/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/rapidfireai/fit/db/db_interface.py b/rapidfireai/fit/db/db_interface.py deleted file mode 100644 index 83fef165..00000000 --- a/rapidfireai/fit/db/db_interface.py +++ /dev/null @@ -1,148 +0,0 @@ -"""Interface for the database.""" - -import functools -import os -import sqlite3 -import time -from collections.abc import Callable -from typing import Any - -from rapidfireai.fit.utils.constants import DBConfig -from rapidfireai.fit.utils.exceptions import DBException - - -class DatabaseInterface: - """Interface for the database.""" - - def __init__(self): - try: - if not os.path.exists(DBConfig.DB_PATH): - path = os.path.dirname(DBConfig.DB_PATH) - os.makedirs(path, exist_ok=True) - print(f"Created directory for database at {path}") - - self.conn: sqlite3.Connection = sqlite3.connect( - DBConfig.DB_PATH, - timeout=DBConfig.CONNECTION_TIMEOUT, - check_same_thread=False, - isolation_level=None, - ) - - # Configure database with all PRAGMA settings - pragma_sql = f""" - PRAGMA cache_size={DBConfig.CACHE_SIZE}; - PRAGMA mmap_size={DBConfig.MMAP_SIZE}; - PRAGMA page_size={DBConfig.PAGE_SIZE}; - PRAGMA busy_timeout={DBConfig.BUSY_TIMEOUT}; - PRAGMA journal_mode=WAL; - PRAGMA synchronous=NORMAL; - PRAGMA temp_store=MEMORY; - PRAGMA foreign_keys=ON; - """ - _ = self.conn.executescript(pragma_sql) - - self.cursor: sqlite3.Cursor = self.conn.cursor() - - except sqlite3.Error as e: - raise DBException(f"Failed to initialize database connection: {e}") from e - except Exception as e: - raise DBException( - f"Unexpected error during database initialization: {e}" - ) from e - - @staticmethod - def retry_on_locked( - max_retries: int = DBConfig.DEFAULT_MAX_RETRIES, - base_delay: float = DBConfig.DEFAULT_BASE_DELAY, - max_delay: float = DBConfig.DEFAULT_MAX_DELAY, - ) -> Callable[[Callable[..., Any]], Callable[..., Any]]: - """Decorator to retry operations when database is locked""" - - def decorator(func: Callable[..., Any]) -> Callable[..., Any]: - @functools.wraps(func) - def wrapper(*args: Any, **kwargs: Any) -> Any: - last_exception = None - for attempt in range(max_retries): - try: - return func(*args, **kwargs) - except sqlite3.OperationalError as e: - if "database is locked" in str(e).lower(): - last_exception = e - if attempt < max_retries - 1: - # Exponential backoff with jitter - delay = min(base_delay * (2**attempt), max_delay) - delay += time.time() % 0.1 # Add small jitter - time.sleep(delay) - continue - # Re-raise if it's not a "database is locked" error - raise - # If we get here, all retries failed - if last_exception: - raise last_exception - else: - raise RuntimeError( - "All retries failed but no exception was captured" - ) - - return wrapper - - return decorator - - def close(self) -> None: - """Close the database connection properly""" - try: - if self.conn: - self.conn.close() - except sqlite3.Error as e: - raise DBException(f"Error closing database connection: {e}") from e - except Exception as e: - raise DBException( - f"Unexpected error closing database connection: {e}" - ) from e - - def optimize_periodically(self) -> None: - """Run periodic optimization - call this occasionally, not on every query""" - try: - _ = self.conn.execute("PRAGMA optimize") - except sqlite3.Error as e: - raise DBException(f"Failed to optimize database: {e}") from e - except Exception as e: - raise DBException( - f"Unexpected error during database optimization: {e}" - ) from e - - def execute( - self, - query: str, - params: dict[str, Any] | tuple[Any, ...] | None = None, - fetch: bool = False, - commit: bool = False, - ) -> list[Any] | tuple[Any] | None: - """Execute a query with automatic retry on database locked errors""" - # Validate that either fetch or commit is True - if not fetch and not commit: - raise ValueError("Either fetch or commit must be True") - - try: - # Execute the query with parameters if provided - if params: - result = self.cursor.execute(query, params) - else: - result = self.cursor.execute(query) - - # Commit the transaction if commit is True - if commit: - self.conn.commit() - - # Return the result if fetch is True - if fetch: - return result.fetchall() - - except sqlite3.Error as e: - raise DBException( - f"Database error executing query '{query[:50]}...': {e}" - ) from e - except Exception as e: - raise DBException( - f"Unexpected error executing query '{query[:50]}...': {e}" - ) from e diff --git a/rapidfireai/fit/db/rf_db.py b/rapidfireai/fit/db/rf_db.py deleted file mode 100644 index 921e8383..00000000 --- a/rapidfireai/fit/db/rf_db.py +++ /dev/null @@ -1,804 +0,0 @@ -"""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.fit.db.db_interface import DatabaseInterface -from rapidfireai.fit.utils.constants import ( - ControllerTask, - ExperimentStatus, - ExperimentTask, - RunEndedBy, - RunSource, - RunStatus, - TaskStatus, - WorkerTask, -) -from rapidfireai.fit.utils.exceptions import DBException -from rapidfireai.fit.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) - else: - try: - cursor = self.db.conn.execute("PRAGMA table_info(runs)") - columns = [column[1] for column in cursor.fetchall()] - if "metric_run_id" not in columns: - self.db.conn.execute( - "ALTER TABLE runs ADD COLUMN metric_run_id TEXT" - ) - self.db.conn.commit() - except sqlite3.Error: - pass - try: - cursor = self.db.conn.execute("PRAGMA table_info(experiments)") - columns = [column[1] for column in cursor.fetchall()] - if "metric_experiment_id" not in columns: - self.db.conn.execute( - "ALTER TABLE experiments ADD COLUMN metric_experiment_id TEXT" - ) - self.db.conn.commit() - except sqlite3.Error: - pass - 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, - metric_experiment_id: str | None, - config_options: dict[str, Any], - ) -> int: - """Create a new experiment""" - query = """ - INSERT INTO experiments (experiment_name, metric_experiment_id, config_options, - status, current_task, error) - VALUES (?, ?, ?, ?, ?, ?) - RETURNING experiment_id - """ - - result = self.db.execute( - query, - ( - experiment_name, - metric_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, metric_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": experiment_details[ - 2 - ], # Return string value, not enum (for JSON serialization) - "error": experiment_details[3], - "metric_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, - metric_run_id: str | None = None, - flattened_config: dict[str, Any] | None = None, - completed_steps: int = 0, - total_steps: int = 0, - num_chunks_visited_curr_epoch: int = 0, - num_epochs_completed: int = 0, - chunk_offset: int = 0, - error: str = "", - source: RunSource | None = None, - ended_by: RunEndedBy | None = None, - warm_started_from: int | None = None, - cloned_from: int | None = None, - estimated_runtime: float = 0.0, - required_workers: int = 0, - ) -> int: - """Create a new run""" - query = """ - INSERT INTO runs (status, metric_run_id, flattened_config, config_leaf, - completed_steps, total_steps, num_chunks_visited_curr_epoch, - num_epochs_completed, chunk_offset, error, source, ended_by, warm_started_from, cloned_from, - estimated_runtime, required_workers) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """ - self.db.execute( - query, - ( - status.value, - metric_run_id, - json.dumps(flattened_config) if flattened_config else "{}", - encode_payload(config_leaf) if config_leaf else "{}", - completed_steps, - total_steps, - num_chunks_visited_curr_epoch, - num_epochs_completed, - chunk_offset, - error, - source.value if source else "", - ended_by.value if ended_by else "", - warm_started_from, - cloned_from, - estimated_runtime, - required_workers, - ), - 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, - metric_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, - num_chunks_visited_curr_epoch: int | None = None, - num_epochs_completed: int | None = None, - chunk_offset: int | None = None, - error: str | None = None, - source: RunSource | None = None, - ended_by: RunEndedBy | None = None, - warm_started_from: int | None = None, - cloned_from: int | None = None, - estimated_runtime: float | None = None, - required_workers: 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, - "metric_run_id": metric_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, - "num_chunks_visited_curr_epoch": num_chunks_visited_curr_epoch, - "num_epochs_completed": num_epochs_completed, - "chunk_offset": chunk_offset, - "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, - "cloned_from": cloned_from, - "estimated_runtime": estimated_runtime, - "required_workers": required_workers, - } - - # 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, metric_run_id, flattened_config, config_leaf, completed_steps, total_steps, - num_chunks_visited_curr_epoch, num_epochs_completed, chunk_offset, error, source, ended_by, - warm_started_from, cloned_from, estimated_runtime, required_workers - 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]), - "metric_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], - "num_chunks_visited_curr_epoch": run_details[6], - "num_epochs_completed": run_details[7], - "chunk_offset": 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], - "cloned_from": run_details[13], - "estimated_runtime": run_details[14], - "required_workers": run_details[15], - } - 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, metric_run_id, flattened_config, config_leaf, completed_steps, total_steps, - num_chunks_visited_curr_epoch, num_epochs_completed, chunk_offset, error, source, ended_by, - warm_started_from, cloned_from, estimated_runtime, required_workers - 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]), - "metric_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], - "num_chunks_visited_curr_epoch": run[7], - "num_epochs_completed": run[8], - "chunk_offset": 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], - "cloned_from": run[14], - "estimated_runtime": run[15], - "required_workers": run[16], - } - 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, metric_run_id, flattened_config, config_leaf, completed_steps, total_steps, - num_chunks_visited_curr_epoch, num_epochs_completed, chunk_offset, error, source, ended_by, - warm_started_from, cloned_from, estimated_runtime, required_workers - 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]), - "metric_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], - "num_chunks_visited_curr_epoch": run[7], - "num_epochs_completed": run[8], - "chunk_offset": 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], - "cloned_from": run[14], - "estimated_runtime": run[15], - "required_workers": run[16], - } - 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") - - def set_estimated_runtime(self, run_id: int, estimated_runtime: float) -> None: - """Set the estimated runtime for a run""" - query = """ - UPDATE runs - SET estimated_runtime = ? - WHERE run_id = ? - """ - self.db.execute(query, (estimated_runtime, run_id), commit=True) - - # 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, - multi_worker_details: dict[str, Any] | None = None, - 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, multi_worker_details, config_options) - VALUES (?, ?, ?, ?, ?, ?, ?) - """ - multi_worker_details_str = ( - json.dumps(multi_worker_details) if multi_worker_details else "{}" - ) - 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, - multi_worker_details_str, - 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, multi_worker_details, 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], - "multi_worker_details": ( - json.loads(task[6]) if task[6] and task[6] != "{}" else {} - ), - "config_options": ( - decode_db_payload(task[7]) - if task[7] and task[7] != "{}" - 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, multi_worker_details, 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] - multi_worker_details = ( - json.loads(task_details[4]) - if task_details[4] and task_details[4] != "{}" - else {} - ) - config_options = ( - decode_db_payload(task_details[5]) - if task_details[5] and task_details[5] != "{}" - else {} - ) - formatted_details = { - "task_id": task_details[0], - "task_type": WorkerTask(task_details[1]), - "run_id": task_details[2], - "chunk_id": task_details[3], - "multi_worker_details": multi_worker_details, - "config_options": ( - config_options - 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 diff --git a/rapidfireai/fit/db/tables.sql b/rapidfireai/fit/db/tables.sql deleted file mode 100644 index 41bd1ed4..00000000 --- a/rapidfireai/fit/db/tables.sql +++ /dev/null @@ -1,68 +0,0 @@ --- Experiments table -CREATE TABLE IF NOT EXISTS experiments ( - experiment_id INTEGER PRIMARY KEY AUTOINCREMENT, - experiment_name TEXT NOT NULL, - metric_experiment_id TEXT, - 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, - metric_run_id TEXT, - flattened_config TEXT DEFAULT '{}', - config_leaf TEXT DEFAULT '{}', - completed_steps INTEGER DEFAULT 0, - total_steps INTEGER DEFAULT 0, - num_chunks_visited_curr_epoch INTEGER DEFAULT 0, - num_epochs_completed INTEGER DEFAULT 0, - chunk_offset INTEGER DEFAULT 0, - error TEXT DEFAULT '', - source TEXT DEFAULT '', - ended_by TEXT DEFAULT '', - warm_started_from BOOLEAN DEFAULT FALSE, - cloned_from INTEGER DEFAULT NULL, - estimated_runtime REAL DEFAULT 0.0, - required_workers INTEGER DEFAULT 0 -); - --- 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 '{}', - multi_worker_details 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) -); \ No newline at end of file diff --git a/rapidfireai/fit/dispatcher/CLAUDE.md b/rapidfireai/fit/dispatcher/CLAUDE.md deleted file mode 100644 index 6d068b2e..00000000 --- a/rapidfireai/fit/dispatcher/CLAUDE.md +++ /dev/null @@ -1,371 +0,0 @@ -# CLAUDE.md - Dispatcher - -This file provides guidance for working with the dispatcher module of RapidFire AI. - -## Overview - -The dispatcher is a Flask-based REST API that provides the communication layer between the web UI (frontend) and the RapidFire backend. It exposes endpoints for viewing experiment status, retrieving run information, and triggering Interactive Control Operations (IC Ops). - -## Files - -### dispatcher.py -**Purpose**: Flask application with REST endpoints for UI communication - -**Key Responsibilities**: -- Serves REST API for frontend dashboard -- Provides endpoints for run queries and experiment status -- Handles IC Ops requests (stop, resume, clone-modify, delete) -- Returns logs for debugging -- Manages CORS for local development - -**Architecture**: -- Flask app with CORS enabled for localhost:8853 (frontend dev server) -- Stateless request handling (reads from database on each request) -- Returns JSON responses -- Error handling with try/catch and HTTP status codes - -**Route Categories**: - -**Health Check**: -- `GET /dispatcher/health-check`: Server health status - -**UI Data Routes**: -- `GET /dispatcher/get-all-runs`: Retrieve all runs for current experiment -- `POST /dispatcher/get-run`: Get single run by ID -- `GET /dispatcher/get-all-experiment-names`: List all experiment names -- `GET /dispatcher/get-running-experiment`: Get currently active experiment - -**Interactive Control Routes**: -- `POST /dispatcher/clone-modify-run`: Clone run with optional modifications -- `POST /dispatcher/stop-run`: Stop active run -- `POST /dispatcher/resume-run`: Resume stopped run -- `POST /dispatcher/delete-run`: Delete run (mark as KILLED) - -**Log Routes**: -- `POST /dispatcher/get-ic-logs`: Get IC Ops logs -- `POST /dispatcher/get-experiment-logs`: Get experiment logs - -**Key Methods**: - -`get_all_runs()`: -- Returns list of all runs with status, metrics, config -- Used by dashboard to display run table -- Includes calculated fields (progress, current_chunk, etc.) - -`clone_modify_run(run_id, config_leaf, warm_start)`: -- Creates IC Ops request in database for clone -- Controller polls and processes request -- Returns new run_id or error - -`stop_run(run_id)`: -- Validates run is in stoppable state (ONGOING) -- Creates IC Ops request in database -- Controller processes asynchronously -- Returns success/error status - -`resume_run(run_id)`: -- Validates run is STOPPED -- Creates IC Ops request in database -- Controller adds run back to scheduler -- Returns success/error status - -`delete_run(run_id)`: -- Marks run as KILLED -- Creates IC Ops request in database -- Controller removes from scheduler -- Returns success/error status - -**Error Handling**: -```python -try: - # ... operation ... - return jsonify(result), 200 -except Exception as e: - logger.error(f"Error: {e}") - return jsonify({"error": str(e)}), 500 -``` - -**CORS Configuration**: -- Allows origins: localhost:8853, localhost -- Required for frontend dev server (separate port from backend) -- Production: frontend built and served from same origin - -### gunicorn.conf.py -**Purpose**: Gunicorn server configuration for production deployment - -**Key Settings**: -- `workers`: Number of worker processes (default: 4) -- `bind`: Host and port (default: 0.0.0.0:8851) -- `timeout`: Request timeout (default: 120s) -- `loglevel`: Log verbosity (default: info) - -**Usage**: -```bash -gunicorn -c rapidfireai/fit/dispatcher/gunicorn.conf.py rapidfireai.fit.dispatcher.dispatcher:app -``` - -**Production Notes**: -- Multiple workers for load balancing -- Timeout prevents hanging requests -- Access logs for monitoring - -## API Endpoints Reference - -### GET /dispatcher/health-check -**Response**: -```json -"Dispatcher is up and running" -``` - -### GET /dispatcher/get-all-runs -**Response**: -```json -[ - { - "run_id": 1, - "run_name": "run_1", - "status": "ONGOING", - "current_chunk": 5, - "current_epoch": 0, - "metrics": "{\"loss\": 0.5, \"accuracy\": 0.9}", - "config_leaf": {...}, - "source": "USER", - "parent_run_id": null, - "warm_start": false - }, - ... -] -``` - -### POST /dispatcher/clone-modify-run -**Request**: -```json -{ - "run_id": 1, - "config_leaf": {"learning_rate": 1e-4}, - "warm_start": true -} -``` -**Response**: -```json -{ - "message": "Clone-modify request created", - "new_run_id": 5 -} -``` - -### POST /dispatcher/stop-run -**Request**: -```json -{ - "run_id": 1 -} -``` -**Response**: -```json -{ - "message": "Stop request created successfully" -} -``` - -### POST /dispatcher/resume-run -**Request**: -```json -{ - "run_id": 1 -} -``` -**Response**: -```json -{ - "message": "Resume request created successfully" -} -``` - -### POST /dispatcher/delete-run -**Request**: -```json -{ - "run_id": 1 -} -``` -**Response**: -```json -{ - "message": "Delete request created successfully" -} -``` - -## Integration with Frontend - -Frontend makes HTTP requests to dispatcher: -```typescript -// Example: Get all runs -const response = await fetch('http://localhost:8851/dispatcher/get-all-runs'); -const runs = await response.json(); - -// Example: Stop run -await fetch('http://localhost:8851/dispatcher/stop-run', { - method: 'POST', - headers: {'Content-Type': 'application/json'}, - body: JSON.stringify({run_id: 5}) -}); -``` - -Frontend polls `get-all-runs` periodically to update dashboard. - -## Integration with Controller - -Dispatcher writes IC Ops requests to database: -```python -# Dispatcher -db.request_stop(run_id) - -# Controller (polling loop) -ic_ops = db.get_ic_ops_request() -for op in ic_ops: - if op['operation'] == 'STOP': - self._handle_stop(op['run_id']) - db.mark_ic_ops_completed(op['ic_id']) -``` - -Asynchronous communication via database (no direct RPC). - -## Running Dispatcher - -**Development**: -```bash -# Via start_dev.sh (starts all services) -./rapidfireai/fit/start_dev.sh start - -# Or manually -python -m flask --app rapidfireai.fit.dispatcher.dispatcher:app run --port 8851 -``` - -**Production** (via start.sh): -```bash -gunicorn -c rapidfireai/fit/dispatcher/gunicorn.conf.py rapidfireai.fit.dispatcher.dispatcher:app -``` - -**Testing**: -```bash -# Health check -curl http://localhost:8851/dispatcher/health-check - -# Get all runs -curl http://localhost:8851/dispatcher/get-all-runs - -# Stop run -curl -X POST http://localhost:8851/dispatcher/stop-run \ - -H "Content-Type: application/json" \ - -d '{"run_id": 1}' -``` - -## Common Patterns - -### Adding New Endpoint - -1. **Add route in `register_routes()`**: -```python -self.app.add_url_rule( - f"{route_prefix}/my-endpoint", - "my_endpoint", - self.my_endpoint, - methods=["POST"] -) -``` - -2. **Implement handler method**: -```python -def my_endpoint(self) -> tuple[Response, int]: - try: - data = request.json - result = self.db.some_operation(data) - return jsonify(result), 200 - except Exception as e: - self._get_logger().error(f"Error: {e}") - return jsonify({"error": str(e)}), 500 -``` - -3. **Update frontend to call endpoint**: -```typescript -await fetch('http://localhost:8851/dispatcher/my-endpoint', { - method: 'POST', - body: JSON.stringify(data) -}); -``` - -### Debugging Dispatcher Issues - -**Check logs**: -```bash -# Dispatcher logs -cat {log_dir}/dispatcher.log - -# Or watch in real-time -tail -f {log_dir}/dispatcher.log -``` - -**Test endpoint directly**: -```bash -curl -X POST http://localhost:8851/dispatcher/stop-run \ - -H "Content-Type: application/json" \ - -d '{"run_id": 1}' \ - -v -``` - -**Check database state**: -```bash -sqlite3 rapidfire.db "SELECT * FROM interactive_control ORDER BY created_at DESC LIMIT 5;" -``` - -**CORS issues**: -- Ensure frontend origin in `CORS_ALLOWED_ORIGINS` -- Check browser console for CORS errors -- Verify preflight OPTIONS requests succeed - -### Error Handling Best Practices - -```python -def endpoint(self) -> tuple[Response, int]: - try: - # Validate input - data = request.json - if not data or 'run_id' not in data: - return jsonify({"error": "run_id required"}), 400 - - # Perform operation - result = self.db.some_operation(data['run_id']) - - # Check result - if not result: - return jsonify({"error": "Operation failed"}), 500 - - return jsonify(result), 200 - except DBException as e: - self._get_logger().error(f"DB error: {e}") - return jsonify({"error": "Database error"}), 500 - except Exception as e: - self._get_logger().error(f"Unexpected error: {e}") - return jsonify({"error": str(e)}), 500 -``` - -## Performance Considerations - -- Dispatcher is stateless (scales horizontally with Gunicorn workers) -- Database is bottleneck for high request volume (SQLite) -- Frontend should throttle polling (e.g., 1-2 second intervals) -- Large run counts may slow `get-all-runs` (consider pagination) - -## Security Notes - -- No authentication (assumes local/trusted network) -- No rate limiting (relies on frontend behavior) -- CORS restricted to localhost (production should tighten) -- Input validation minimal (assumes trusted clients) - -For production deployment, consider adding: -- API authentication (tokens, JWT) -- Rate limiting (Flask-Limiter) -- Input validation (marshmallow, pydantic) -- HTTPS (reverse proxy like nginx) diff --git a/rapidfireai/fit/dispatcher/dispatcher.py b/rapidfireai/fit/dispatcher/dispatcher.py deleted file mode 100644 index 05e1bc90..00000000 --- a/rapidfireai/fit/dispatcher/dispatcher.py +++ /dev/null @@ -1,493 +0,0 @@ -"""This module contains functions for the dispatcher module.""" - -import os -import traceback -from logging import Logger -from typing import Any - -from flask import Flask, Response, jsonify, request -from flask_cors import CORS - -from rapidfireai.fit.db.rf_db import RfDb -from rapidfireai.utils.constants import DispatcherConfig, FrontendConfig, MLflowConfig, RF_LOG_FILENAME, RF_LOG_PATH -from rapidfireai.utils.dispatcher_utils import check_experiment_running -from rapidfireai.fit.utils.constants import ControllerTask -from rapidfireai.fit.utils.exceptions import DispatcherException -from rapidfireai.fit.utils.logging import RFLogger - -CORS_ALLOWED_ORIGINS = ["http://localhost", DispatcherConfig.URL, MLflowConfig.URL, FrontendConfig.URL] - - -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 - # Allow all origins for local development (dispatcher runs on localhost) - # This is safe since the API is not exposed to the internet - _ = CORS(self.app, resources={r"/*": {"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"], - ) - self.app.add_url_rule( - f"{route_prefix}/is-experiment-running", - "is_experiment_running", - self.is_experiment_running, - methods=["POST"], - ) - - # 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, - "metric_run_id": result["metric_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, - "metric_run_id": result["metric_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 - - def is_experiment_running(self) -> tuple[Response, int]: - """Check if a specific experiment is currently running. - - Request body: - { - "experiment_name": str - The experiment name to check - } - - Returns: - { - "is_running": bool - True if the experiment is currently running - } - """ - if request.method == "OPTIONS": - return jsonify({}), 200 - - try: - data = request.get_json() - if not data or "experiment_name" not in data: - return jsonify({"error": "experiment_name is required"}), 400 - - is_running = check_experiment_running(self.db, data["experiment_name"]) - return jsonify({"is_running": is_running}), 200 - except Exception: - # If anything fails, assume experiment is not running (safer to disable button) - return jsonify({"is_running": False}), 200 - - # 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 = os.path.join(RF_LOG_PATH, experiment_name, RF_LOG_FILENAME) - - # Check if the log file exists - if not os.path.exists(log_file_path): - 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, encoding="utf-8") 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 = os.path.join(RF_LOG_PATH, experiment_name, RF_LOG_FILENAME) - - experiment_logs = [] - with open(log_file_path, encoding="utf-8") 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) diff --git a/rapidfireai/fit/ml/callbacks.py b/rapidfireai/fit/ml/callbacks.py index 905f8027..5783d83c 100644 --- a/rapidfireai/fit/ml/callbacks.py +++ b/rapidfireai/fit/ml/callbacks.py @@ -13,7 +13,6 @@ from transformers.trainer_utils import IntervalStrategy, SaveStrategy from rapidfireai.fit.ml.checkpoint_utils import flush_cuda_cache, purge_model_kv_caches -from rapidfireai.fit.utils.logging import RFLogger class GenerationMetricsCallback(TrainerCallback): @@ -418,6 +417,8 @@ def __init__( completed_steps: int = 0, chunk_id: int = 0, num_epochs_completed: int = 0, + db=None, + run_id: int | None = None, ): self.metric_logger = metric_logger self.metric_run_id = metric_run_id @@ -428,6 +429,20 @@ def __init__( ] self.chunk_id = chunk_id self.num_epochs_completed = num_epochs_completed + self.db = db + self.run_id = run_id + + def _is_run_deleted(self) -> bool: + """Check if the run has been deleted via Interactive Control.""" + if self.db is None or self.run_id is None: + return False + try: + from rapidfireai.utils.constants import RunStatus + + run = self.db.get_run(self.run_id) + return run is not None and run.get("status") == RunStatus.DELETED.value + except Exception: + return False def on_log( self, @@ -438,23 +453,30 @@ def on_log( **kwargs, ): """Called when the trainer logs metrics""" - if logs is not None: - step = self.completed_steps + state.global_step - for key, value in logs.items(): - if isinstance(value, (int, float)) and key not in self.excluded_keys: - try: - if self.metric_logger: - self.metric_logger.log_metric( - self.metric_run_id, - key, - value, - step=step, - ) - except Exception as e: - print( - f"Warning: Failed to log metric {key} to tracking backend: {e}" + if logs is None: + return + + if self._is_run_deleted(): + control.should_training_stop = True + return + + step = self.completed_steps + state.global_step + for key, value in logs.items(): + if isinstance(value, (int, float)) and key not in self.excluded_keys: + try: + if self.metric_logger: + self.metric_logger.log_metric( + self.metric_run_id, + key, + value, + step=step, ) - if "eval_loss" not in logs and "train_runtime" not in logs: + except Exception as e: + print( + f"Warning: Failed to log metric {key} to tracking backend: {e}" + ) + if "eval_loss" not in logs and "train_runtime" not in logs: + try: if self.metric_logger: self.metric_logger.log_metric( self.metric_run_id, @@ -468,6 +490,10 @@ def on_log( self.num_epochs_completed, step=step, ) + except Exception as e: + print( + f"Warning: Failed to log chunk/epoch metrics to tracking backend: {e}" + ) class LogLevelCallback(TrainerCallback): diff --git a/rapidfireai/fit/ml/checkpoint_utils.py b/rapidfireai/fit/ml/checkpoint_utils.py index 5e2d8134..1514fee9 100644 --- a/rapidfireai/fit/ml/checkpoint_utils.py +++ b/rapidfireai/fit/ml/checkpoint_utils.py @@ -22,7 +22,7 @@ from transformers import AutoTokenizer from trl import DPOTrainer, GRPOTrainer, SFTTrainer -from rapidfireai.fit.utils.constants import SHMObjectType +from rapidfireai.utils.constants import SHMObjectType from rapidfireai.fit.utils.datapaths import DataPath from rapidfireai.fit.utils.shm_manager import SharedMemoryManager from rapidfireai.fit.utils.trainer_config import TrainerConfig diff --git a/rapidfireai/fit/ml/trainer.py b/rapidfireai/fit/ml/trainer.py index bec00265..b104d362 100644 --- a/rapidfireai/fit/ml/trainer.py +++ b/rapidfireai/fit/ml/trainer.py @@ -1,6 +1,5 @@ import math import os -import warnings import torch from peft import LoraConfig, get_peft_model_state_dict, set_peft_model_state_dict @@ -22,7 +21,7 @@ restore_trainer_from_disk, restore_trainer_from_shared_memory, ) -from rapidfireai.fit.utils.constants import SHMObjectType +from rapidfireai.utils.constants import SHMObjectType from rapidfireai.fit.utils.datapaths import DataPath from rapidfireai.fit.utils.shm_manager import SharedMemoryManager from rapidfireai.fit.utils.trainer_config import TrainerConfig @@ -84,7 +83,7 @@ def restore_optimizer_state(self): self.optimizer.load_state_dict(device_optimizer_state) delattr(self, "_pending_optimizer_state") except Exception as e: - print(f"Warning: Error restoring FSDP scheduler state: {e}") + print(f"Warning: Error restoring optimizer state: {e}") def restore_fsdp_optimizer_state(self, rank=0): from torch.distributed.fsdp.fully_sharded_data_parallel import ( @@ -368,6 +367,7 @@ def create_trainer_instance( metric_logger=None, chunk_id: int = 0, use_fsdp: bool = False, + db=None, ) -> tuple[SFTTrainer | DPOTrainer | GRPOTrainer | None, str]: """ Create a trainer instance with proper state restoration. @@ -446,6 +446,7 @@ def create_trainer_instance( formatting_func, global_step_args, use_fsdp=use_fsdp, + db=db, ) ) @@ -533,9 +534,13 @@ def _configure_training_args( if "save_steps" in training_args: training_args.pop("save_steps") - # Configure distributed training arguments for FSDP + # Configure distributed training arguments for FSDP. + # Each Ray actor sees only its assigned GPU as cuda:0, so the HF Trainer device + # index is always 0. trainer_config.local_rank holds the FSDP group rank (used + # for conditional logic like "only rank 0 saves checkpoints") and must not be + # used as a device index here. if use_fsdp: - training_args["local_rank"] = trainer_config.local_rank + training_args["local_rank"] = 0 # training_args["dataloader_num_workers"] = trainer_config.world_size # Avoid multiprocessing issues with FSDP training_args["dataloader_pin_memory"] = False training_args["remove_unused_columns"] = False # FSDP requires this @@ -694,6 +699,7 @@ def _setup_callbacks( formatting_func, global_step_args, use_fsdp=False, + db=None, ): """Setup callbacks for the trainer.""" callbacks = [] @@ -705,6 +711,8 @@ def _setup_callbacks( completed_steps=trainer_config.completed_steps, chunk_id=chunk_id, num_epochs_completed=trainer_config.num_epochs_completed, + db=db, + run_id=trainer_config.run_id, ) callbacks.append(metric_callback) diff --git a/rapidfireai/fit/utils/constants.py b/rapidfireai/fit/utils/constants.py deleted file mode 100644 index 70c8b6f5..00000000 --- a/rapidfireai/fit/utils/constants.py +++ /dev/null @@ -1,130 +0,0 @@ -import os -from enum import Enum -from rapidfireai.utils.constants import RF_TRAINING_LOG_FILENAME, RF_DB_PATH, ExperimentStatus - -# Backwards compatibility: Keep constant but it will be stale if env var changes after import - -# Shared Memory Constants -SHM_WARN_THRESHOLD = 80 -SHM_MIN_FREE_SPACE = 1.0 -USE_SHARED_MEMORY = True - -# Logging Constants -TRAINING_LOG_FILENAME = RF_TRAINING_LOG_FILENAME - - -class LogType(Enum): - """Enum class for log types""" - - RF_LOG = "rf_log" - TRAINING_LOG = "training_log" - - -# Database Constants -class DBConfig: - """Class to manage the database configuration for SQLite""" - - # Use user's home directory for database path - - DB_PATH: str = os.path.join( - RF_DB_PATH, "rapidfire_fit.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" - - -# Note: ExperimentStatus is imported from rapidfireai.utils.constants (shared) - - -# 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" diff --git a/rapidfireai/fit/utils/datapaths.py b/rapidfireai/fit/utils/datapaths.py index b435d097..c91633ce 100644 --- a/rapidfireai/fit/utils/datapaths.py +++ b/rapidfireai/fit/utils/datapaths.py @@ -4,7 +4,7 @@ from pathlib import Path -from rapidfireai.fit.utils.exceptions import DataPathException +from rapidfireai.utils.exceptions import DataPathException from rapidfireai.utils.os_utils import mkdir_p diff --git a/rapidfireai/fit/utils/exceptions.py b/rapidfireai/fit/utils/exceptions.py deleted file mode 100644 index 64133640..00000000 --- a/rapidfireai/fit/utils/exceptions.py +++ /dev/null @@ -1,82 +0,0 @@ -""" -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) diff --git a/rapidfireai/fit/utils/logging.py b/rapidfireai/fit/utils/logging.py deleted file mode 100644 index 303760b7..00000000 --- a/rapidfireai/fit/utils/logging.py +++ /dev/null @@ -1,87 +0,0 @@ -import os -import threading -from abc import ABC, abstractmethod - -from loguru import logger - -from rapidfireai.fit.db.rf_db import RfDb -from rapidfireai.utils.constants import RF_LOG_FILENAME, RF_LOG_PATH -from rapidfireai.fit.utils.constants import TRAINING_LOG_FILENAME, LogType - - -class BaseRFLogger(ABC): - """Base class for RapidFire loggers""" - - _experiment_name = "" - _initialized_loggers: dict[str, bool] = {} - _lock = threading.Lock() - - def __init__(self, level: str = "DEBUG"): - try: - db = RfDb() - experiment_name = db.get_running_experiment()["experiment_name"] - except Exception: - experiment_name = "no_active_experiment" - log_file_path = self.get_log_file_path(experiment_name) - - with BaseRFLogger._lock: - # Reset loggers if experiment changed - if experiment_name != BaseRFLogger._experiment_name: - BaseRFLogger._experiment_name = experiment_name - BaseRFLogger._initialized_loggers = {} - logger.remove() - - # Each process gets its own handler per logger type - logger_type = self.get_logger_type() - logger_key = f"{logger_type.value}_{experiment_name}_{os.getpid()}" - if logger_key not in BaseRFLogger._initialized_loggers: - logger.add( - log_file_path, - format="{time:YYYY-MM-DD HH:mm:ss} | " - + "{extra[experiment_name]} | " - + "{extra[logger_name]} | {level} | " - + "{file}:{line} | {message}", - level=level.upper(), - enqueue=True, - filter=lambda record, logger_type=logger_type: (record["extra"].get("logger_type") == logger_type), - ) - BaseRFLogger._initialized_loggers[logger_key] = True - - @abstractmethod - def get_log_file_path(self, experiment_name: str): - """Get the log file path for this logger type""" - pass - - @abstractmethod - def get_logger_type(self) -> LogType: - """Get the logger type identifier""" - pass - - def create_logger(self, name: str): - """Create a configured logger instance""" - return logger.bind( - logger_name=name, - experiment_name=BaseRFLogger._experiment_name, - logger_type=self.get_logger_type(), - pid=os.getpid(), - ) - - -class RFLogger(BaseRFLogger): - """Standard RapidFire logger""" - - def get_log_file_path(self, experiment_name: str): - return os.path.join(RF_LOG_PATH, experiment_name, RF_LOG_FILENAME) - - def get_logger_type(self) -> LogType: - return LogType.RF_LOG - - -class TrainingLogger(BaseRFLogger): - """Training-specific logger""" - - def get_log_file_path(self, experiment_name: str): - return os.path.join(RF_LOG_PATH, experiment_name, TRAINING_LOG_FILENAME) - - def get_logger_type(self) -> LogType: - return LogType.TRAINING_LOG diff --git a/rapidfireai/fit/utils/serialize.py b/rapidfireai/fit/utils/serialize.py deleted file mode 100644 index 34f4e0d6..00000000 --- a/rapidfireai/fit/utils/serialize.py +++ /dev/null @@ -1,15 +0,0 @@ -"""This module contains the serialize functions for complexobjects.""" - -import base64 - -import dill - - -def encode_payload(payload: object) -> str: - """Encode the payload for the database""" - return base64.b64encode(dill.dumps(payload)).decode("utf-8") - - -def decode_db_payload(payload: str) -> object: - """Decode the payload from the database""" - return dill.loads(base64.b64decode(payload)) diff --git a/rapidfireai/fit/utils/shm_manager.py b/rapidfireai/fit/utils/shm_manager.py index 7cb546cb..27293c95 100644 --- a/rapidfireai/fit/utils/shm_manager.py +++ b/rapidfireai/fit/utils/shm_manager.py @@ -1,17 +1,89 @@ +""" +Shared Memory Manager for PyTorch models and checkpoints. + +This module provides: +- RegistryActor: A Ray actor that serves as a shared registry for model metadata +- SharedMemoryManager: Manages PyTorch tensors in /dev/shm with zero-copy semantics + +The actual tensor data is stored in /dev/shm via torch.Tensor.share_memory_(). +The RegistryActor provides coordination between workers (which models exist, locks, etc.). +""" + import copy import gc import threading -from multiprocessing import Lock, Manager +from typing import Any +import ray import torch -from rapidfireai.fit.utils.constants import ( - SHM_MIN_FREE_SPACE, - SHM_WARN_THRESHOLD, - SHMObjectType, -) -from rapidfireai.fit.utils.exceptions import InsufficientSharedMemoryException -from rapidfireai.fit.utils.logging import RFLogger +from rapidfireai.utils.constants import SHMObjectType +from rapidfireai.utils.exceptions import InsufficientSharedMemoryException +from rapidfireai.utils.logging import RFLogger + +SHM_WARN_THRESHOLD = 80 +SHM_MIN_FREE_SPACE = 1.0 +USE_SHARED_MEMORY = True + + +@ray.remote +class RegistryActor: + """ + Ray Actor that serves as a shared registry for model metadata. + + This replaces multiprocessing.Manager().dict() with a Ray-native solution. + Ray actors are single-threaded, so method calls are automatically serialized + (no explicit locking needed). + + The registry stores model objects whose tensors are in /dev/shm. + When these objects are returned to callers, the tensors remain zero-copy + because PyTorch preserves shared memory handles during serialization. + """ + + def __init__(self): + self._data: dict[str, Any] = {} + + def get(self, key: str) -> Any: + """Get a value from the registry.""" + return self._data.get(key) + + def set(self, key: str, value: Any) -> None: + """Set a value in the registry.""" + self._data[key] = value + + def delete(self, key: str) -> bool: + """Delete a key from the registry. Returns True if key existed.""" + if key in self._data: + del self._data[key] + return True + return False + + def contains(self, key: str) -> bool: + """Check if a key exists in the registry.""" + return key in self._data + + def keys(self) -> list[str]: + """Get all keys in the registry.""" + return list(self._data.keys()) + + def set_entry_field(self, key: str, field: SHMObjectType, value: Any) -> None: + """Set a specific field in a registry entry, creating entry if needed.""" + if key not in self._data: + self._data[key] = {} + self._data[key][field] = value + + def delete_entry_field(self, key: str, field: SHMObjectType) -> bool: + """Delete a specific field from a registry entry.""" + if key in self._data and field in self._data[key]: + del self._data[key][field] + return True + return False + + def get_or_create_entry(self, key: str, default: dict) -> dict: + """Get an entry, creating it with default value if it doesn't exist.""" + if key not in self._data: + self._data[key] = default + return self._data[key] def _get_shm_usage(): @@ -120,27 +192,41 @@ def _verify_sufficient_ref_state_dict_size(ref_state_dict: dict, logger: RFLogge class SharedMemoryManager: - """Manages PyTorch models and checkpoints in shared memory across multiple processes.""" - - def __init__(self, name: str, registry=None, multiprocess_lock=None): - """Initialize the shared memory manager with process-safe registry and locks""" - # initialize registry - if registry is None: - self._manager = Manager() - self._registry = self._manager.dict() - else: - self._registry = registry + """ + Manages PyTorch models and checkpoints in shared memory across multiple processes. + + This class handles: + - Moving tensors to /dev/shm via torch.Tensor.share_memory_() + - BitsAndBytes quantization state handling + - Coordinating with RegistryActor for model metadata - # initialize multiprocess lock - if multiprocess_lock is None: - self._process_lock = self._manager.Lock() + The actual tensor data is zero-copy shared between all workers on the same machine. + """ + + def __init__(self, name: str, registry_actor=None): + """ + Initialize the shared memory manager. + + Args: + name: Name for logging purposes + registry_actor: Ray actor handle for RegistryActor. If None, creates a new one. + """ + # Initialize or use provided registry actor + if registry_actor is None: + self._registry_actor = RegistryActor.remote() + self._owns_registry = True else: - self._process_lock = multiprocess_lock + self._registry_actor = registry_actor + self._owns_registry = False - # initialize thread lock for operations within a single process + # Thread lock for operations within a single process self._thread_lock = threading.Lock() - self.logger = RFLogger().create_logger(name) + self.logger = RFLogger().get_logger(name) + + def get_registry_actor(self): + """Get the registry actor handle (for passing to workers).""" + return self._registry_actor # shared memory operations def _safe_tensor_to_shared_memory( @@ -301,34 +387,18 @@ def _save_full_model( is_fsdp=False, ): """Save the full model in shared memory. model_id can be either run_id or name of a base model""" - with self._process_lock if self._process_lock else self._thread_lock: - if ( - model_id in self._registry - and model_object_type != SHMObjectType.FULL_MODEL - ): - self.logger.debug( - f"Model {model_id} already exists in shared memory. Skipping save." - ) - if ( - model_id in self._registry - and model_object_type != SHMObjectType.FULL_MODEL - ): - self.logger.debug( - f"Model {model_id} already exists in shared memory. Skipping save." - ) + with self._thread_lock: + if ray.get(self._registry_actor.contains.remote(model_id)) and model_object_type != SHMObjectType.FULL_MODEL: + self.logger.debug(f"Model {model_id} already exists in shared memory. Skipping save.") return # verify sufficient shared memory space before saving model _verify_sufficient_model_size( model_data[model_object_type.value], self.logger ) - _verify_sufficient_model_size( - model_data[model_object_type.value], self.logger - ) - # create model entry in registry - if model_id not in self._registry: - self._registry[model_id] = {model_object_type: {}} + # create model entry in registry if needed + ray.get(self._registry_actor.get_or_create_entry.remote(model_id, {model_object_type: {}})) # move model to shared memory model_cpu = model_data[model_object_type.value] @@ -341,9 +411,9 @@ def _save_full_model( "tokenizer": tokenizer, "bnb_modules": self._move_tensors_to_shared_memory(bnb_modules), } - model_entry = dict(self._registry[model_id]) - model_entry[model_object_type] = shared_model - self._registry[model_id] = model_entry + + # Update the entry + ray.get(self._registry_actor.set_entry_field.remote(model_id, model_object_type, shared_model)) self.logger.debug(f"Saved {model_object_type.value} for run {model_id}") @@ -353,26 +423,28 @@ def _save_ref_state_dict(self, model_id: str, ref_state_dict: dict): # verify sufficient shared memory space before saving ref_state_dict _verify_sufficient_ref_state_dict_size(ref_state_dict, self.logger) - # create model entry in registry - if model_id not in self._registry: - self._registry[model_id] = {SHMObjectType.REF_STATE_DICT: {}} + # create model entry in registry if needed + ray.get(self._registry_actor.get_or_create_entry.remote(model_id, {SHMObjectType.REF_STATE_DICT: {}})) # move ref_state_dict to shared memory shared_ref_state_dict = self._move_tensors_to_shared_memory(ref_state_dict) - model_entry = dict(self._registry[model_id]) - model_entry[SHMObjectType.REF_STATE_DICT] = shared_ref_state_dict - self._registry[model_id] = model_entry + + # Update the entry + ray.get(self._registry_actor.set_entry_field.remote(model_id, SHMObjectType.REF_STATE_DICT, shared_ref_state_dict)) self.logger.debug(f"Saved ref_state_dict for {model_id}") def _update_checkpoints(self, model_id: str, checkpoint_updates: dict): """Update checkpoints in-place when possible, add new keys when needed.""" with self._thread_lock: - # create model entry in registry - if model_id not in self._registry: - self._registry[model_id] = {SHMObjectType.CHECKPOINTS: {}} + # Get or create entry with checkpoints + ray.get(self._registry_actor.get_or_create_entry.remote(model_id, {SHMObjectType.CHECKPOINTS: {}})) + + # Get current checkpoints + model_entry = ray.get(self._registry_actor.get.remote(model_id)) + if model_entry is None: + model_entry = {} - model_entry = self._registry[model_id] if SHMObjectType.CHECKPOINTS not in model_entry: model_entry[SHMObjectType.CHECKPOINTS] = {} current_checkpoints = model_entry[SHMObjectType.CHECKPOINTS] @@ -437,22 +509,17 @@ def update_nested_dict(current_dict, updates_dict, path=""): # Update the checkpoints update_nested_dict(current_checkpoints, checkpoint_updates) - # Update the registry entry to ensure Manager sees changes - updated_entry = dict(model_entry) - updated_entry[SHMObjectType.CHECKPOINTS] = current_checkpoints - self._registry[model_id] = updated_entry + # Update the registry + model_entry[SHMObjectType.CHECKPOINTS] = current_checkpoints + ray.get(self._registry_actor.set.remote(model_id, model_entry)) self.logger.debug( f"Checkpoint update:{updates_made['in_place']} in-place, {updates_made['new_keys']} new" ) - def get_shm_objects(self) -> tuple[dict, Lock]: - """Get the shared registry and process lock""" - return self._registry, self._process_lock - def load_model_object(self, model_id: str, model_object_type: SHMObjectType): """Load a model object from shared memory.""" - model_entry = self._registry.get(model_id) + model_entry = ray.get(self._registry_actor.get.remote(model_id)) if model_entry is None: self.logger.warning(f"Model {model_id} not found in shared memory") return None @@ -481,76 +548,39 @@ def save_model_object( elif model_object_type == SHMObjectType.CHECKPOINTS: self._update_checkpoints(model_id, model_object) - def delete_model_object( - self, model_id: str, base_model_name: str | None = None, is_fsdp: bool = False - ): - """Delete model object from shared memory registry and clean up resources. - For FSDP, base model is not stored in shared memory - """ - with self._process_lock if self._process_lock else self._thread_lock: - if model_id not in self._registry: - self.logger.warning( - f"Model '{model_id}' not found in shared memory during delete" - ) + def delete_model_object(self, model_id: str, base_model_name: str | None = None, is_fsdp: bool = False): + """Delete model object from shared memory registry and clean up resources.""" + with self._thread_lock: + if not ray.get(self._registry_actor.contains.remote(model_id)): + self.logger.warning(f"Model '{model_id}' not found in shared memory during delete") return + # Get the model entry + model_entry = ray.get(self._registry_actor.get.remote(model_id)) + # remove checkpoints - # TODO: add code to save to disk before deleting - if ( - SHMObjectType.CHECKPOINTS in self._registry[model_id] - and self._registry[model_id][SHMObjectType.CHECKPOINTS] - ): - del self._registry[model_id][SHMObjectType.CHECKPOINTS] - self.logger.debug( - f"Deleted checkpoints for model {model_id} from shared memory" - ) - - # remove full_model - # TODO: add code to save to disk before deleting - if ( - SHMObjectType.FULL_MODEL in self._registry[model_id] - and self._registry[model_id][SHMObjectType.FULL_MODEL] - ): - del self._registry[model_id][SHMObjectType.FULL_MODEL] - self.logger.debug( - f"Deleted full_model for model {model_id} from shared memory" - ) - self.logger.debug( - f"Deleted full_model for model {model_id} from shared memory" - ) - - # remove ref_state_dict - if ( - SHMObjectType.REF_STATE_DICT in self._registry[model_id] - and self._registry[model_id][SHMObjectType.REF_STATE_DICT] - ): - del self._registry[model_id][SHMObjectType.REF_STATE_DICT] - self.logger.debug( - f"Deleted ref_state_dict for model {model_id} from shared memory" - ) - - # remove ref_full_model - if ( - SHMObjectType.REF_FULL_MODEL in self._registry[model_id] - and self._registry[model_id][SHMObjectType.REF_FULL_MODEL] - ): - del self._registry[model_id][SHMObjectType.REF_FULL_MODEL] - self.logger.debug( - f"Deleted ref_full_model for model {model_id} from shared memory" - ) - - # remove base model from shm (not present for FSDP - base model reloaded from HuggingFace) - if not is_fsdp and base_model_name and base_model_name in self._registry: - del self._registry[base_model_name] - self.logger.debug( - f"Deleted base_model for model {model_id} from shared memory" - ) - - # remove registry entry - del self._registry[model_id] - self.logger.debug( - f"Deleted model registry entry for {model_id} from shared memory" - ) + if model_entry and SHMObjectType.CHECKPOINTS in model_entry and model_entry[SHMObjectType.CHECKPOINTS]: + ray.get(self._registry_actor.delete_entry_field.remote(model_id, SHMObjectType.CHECKPOINTS)) + self.logger.debug(f"Deleted checkpoints for model {model_id} from shared memory") + + if model_entry and SHMObjectType.FULL_MODEL in model_entry and model_entry[SHMObjectType.FULL_MODEL]: + ray.get(self._registry_actor.delete_entry_field.remote(model_id, SHMObjectType.FULL_MODEL)) + self.logger.debug(f"Deleted full_model for model {model_id} from shared memory") + + if model_entry and SHMObjectType.REF_STATE_DICT in model_entry and model_entry[SHMObjectType.REF_STATE_DICT]: + ray.get(self._registry_actor.delete_entry_field.remote(model_id, SHMObjectType.REF_STATE_DICT)) + self.logger.debug(f"Deleted ref_state_dict for model {model_id} from shared memory") + + if model_entry and SHMObjectType.REF_FULL_MODEL in model_entry and model_entry[SHMObjectType.REF_FULL_MODEL]: + ray.get(self._registry_actor.delete_entry_field.remote(model_id, SHMObjectType.REF_FULL_MODEL)) + self.logger.debug(f"Deleted ref_full_model for model {model_id} from shared memory") + + if not is_fsdp and base_model_name: + ray.get(self._registry_actor.delete.remote(base_model_name)) + self.logger.debug(f"Deleted base_model for model {model_id} from shared memory") + + ray.get(self._registry_actor.delete.remote(model_id)) + self.logger.debug(f"Deleted model registry entry for {model_id} from shared memory") # Force garbage collection gc.collect() @@ -564,44 +594,33 @@ def create_warm_start_checkpoint( ): """Copy warm start checkpoint from run_id to cloned_from.""" with self._thread_lock: - if warm_started_from not in self._registry: + if not ray.get(self._registry_actor.contains.remote(warm_started_from)): raise KeyError(f"Run '{warm_started_from}' not found in shared memory") - # create model entry in registry - if model_id not in self._registry: - self._registry[model_id] = { - SHMObjectType.FULL_MODEL: {}, - SHMObjectType.REF_STATE_DICT: {}, - SHMObjectType.CHECKPOINTS: {}, - } - - # copy full_model, ref_state_dict, and checkpoints from cloned_from to model_id - model_entry = dict(self._registry[model_id]) - if SHMObjectType.FULL_MODEL in self._registry[warm_started_from]: - model_entry[SHMObjectType.FULL_MODEL] = copy.deepcopy( - dict(self._registry[warm_started_from])[SHMObjectType.FULL_MODEL] - ) - if SHMObjectType.REF_STATE_DICT in self._registry[warm_started_from]: - model_entry[SHMObjectType.REF_STATE_DICT] = copy.deepcopy( - dict(self._registry[warm_started_from])[ - SHMObjectType.REF_STATE_DICT - ] - ) - if SHMObjectType.CHECKPOINTS in self._registry[warm_started_from]: - model_entry[SHMObjectType.CHECKPOINTS] = copy.deepcopy( - dict(self._registry[warm_started_from])[SHMObjectType.CHECKPOINTS] - ) - self._registry[model_id] = model_entry - self.logger.debug( - f"Copied warm start checkpoint from run {warm_started_from} to run {model_id}" - ) + # Get source entry + source_entry = ray.get(self._registry_actor.get.remote(warm_started_from)) + + new_entry = { + SHMObjectType.FULL_MODEL: {}, + SHMObjectType.REF_STATE_DICT: {}, + SHMObjectType.CHECKPOINTS: {}, + } + + if source_entry: + if SHMObjectType.FULL_MODEL in source_entry: + new_entry[SHMObjectType.FULL_MODEL] = copy.deepcopy(source_entry[SHMObjectType.FULL_MODEL]) + if SHMObjectType.REF_STATE_DICT in source_entry: + new_entry[SHMObjectType.REF_STATE_DICT] = copy.deepcopy(source_entry[SHMObjectType.REF_STATE_DICT]) + if SHMObjectType.CHECKPOINTS in source_entry: + new_entry[SHMObjectType.CHECKPOINTS] = copy.deepcopy(source_entry[SHMObjectType.CHECKPOINTS]) + + ray.get(self._registry_actor.set.remote(model_id, new_entry)) + self.logger.debug(f"Copied warm start checkpoint from run {warm_started_from} to run {model_id}") def list_models(self): """Get list of all model IDs currently in shared memory.""" - with self._process_lock if self._process_lock else self._thread_lock: - return list(self._registry.keys()) + return ray.get(self._registry_actor.keys.remote()) def model_exists(self, model_id: str): """Check if a model exists in shared memory.""" - with self._process_lock if self._process_lock else self._thread_lock: - return model_id in self._registry + return ray.get(self._registry_actor.contains.remote(model_id)) diff --git a/rapidfireai/fit/utils/worker_manager.py b/rapidfireai/fit/utils/worker_manager.py deleted file mode 100644 index fa19bf27..00000000 --- a/rapidfireai/fit/utils/worker_manager.py +++ /dev/null @@ -1,240 +0,0 @@ -""" -WorkerManager class for managing worker processes with PyTorch multiprocessing. -""" - -import os -import signal -import sys -import threading -import time -from multiprocessing.managers import DictProxy -from multiprocessing.synchronize import Event as EventType -from multiprocessing.synchronize import Lock - -from torch.multiprocessing import Event, Process - -from rapidfireai.fit.backend.worker import Worker -from rapidfireai.fit.utils.logging import RFLogger - - -def worker_process_target( - worker_id: int, - model_registry: DictProxy, - process_lock: Lock, - shutdown_event: EventType, -): - """ - Target function that runs in each worker process. - Creates Worker instance inside the process to avoid pickling issues. - """ - - # Create worker instance inside the process (avoids pickling) - worker = Worker(worker_id, model_registry, process_lock, shutdown_event) - - # Set up signal handlers for graceful shutdown - def signal_handler(signum, frame): - if signum == signal.SIGTERM: - worker.logger.debug(f"Worker {worker.worker_id} received SIGTERM") - shutdown_event.set() - # Don't exit immediately, let the main loop handle shutdown - - signal.signal(signal.SIGTERM, signal_handler) - - try: - worker.serve_forever() - except KeyboardInterrupt: - worker.logger.debug(f"Worker {worker.worker_id} interrupted") - except Exception as e: - worker.logger.debug(f"Worker {worker.worker_id} crashed: {e}") - finally: - worker.logger.info(f"Worker {worker.worker_id} process exiting") - # Ensure proper cleanup before exit - try: - worker.shutdown() - except Exception as e: - worker.logger.debug(f"Error during worker shutdown: {e}") - # Use sys.exit instead of os._exit for better cleanup - sys.exit(0) - - -class WorkerManager: - """ - WorkerManager class for managing worker processes. - """ - - def __init__( - self, - num_workers: int, - model_registry: DictProxy | None = None, - process_lock: Lock | None = None, - parent_check_interval: float = 1.0, - ): - self.num_workers = num_workers - self.parent_check_interval = parent_check_interval - self.parent_pid = os.getppid() - self.process_group_id: int | None = None - self.processes: list[Process] = [] - self.shutdown_events: list[EventType] = [] - self.shutdown_event = threading.Event() - self.monitor_thread: threading.Thread | None = None - self.model_registry: DictProxy = model_registry - self.process_lock: Lock = process_lock - - self.logger = RFLogger().create_logger("worker_manager") - - def _parent_monitor(self): - """Monitor parent process and kill process group if parent dies""" - while not self.shutdown_event.is_set(): - try: - # Check if parent process is still alive - try: - os.getpgid(self.parent_pid) - except PermissionError: - # Fallback for restricted environments (e.g., Colab) - # Use psutil to check if parent process exists - import psutil - - if not psutil.pid_exists(self.parent_pid): - raise ProcessLookupError("Parent process no longer exists") - time.sleep(self.parent_check_interval) - except ProcessLookupError: - self.logger.debug( - f"Parent process {self.parent_pid} died, shutting down workers..." - ) - self.shutdown() - break - except Exception as e: - self.logger.error(f"Parent monitor error: {e}") - time.sleep(self.parent_check_interval) - - def create_workers(self) -> list[int]: - """ - Create worker processes and return worker IDs. - Returns list of worker IDs to avoid pickling Worker objects. - """ - self.logger.debug(f"Creating {self.num_workers} worker processes...") - - # Create new process group (may not be permitted in restricted environments like Colab) - try: - os.setpgrp() - self.process_group_id = os.getpgrp() - self.logger.debug( - f"Starting worker processes in process group {self.process_group_id}" - ) - except PermissionError: - self.logger.debug( - "Cannot create process group (restricted environment) - will use individual process termination" - ) - self.process_group_id = None - - worker_ids = [] - - # Start each worker process - for i in range(self.num_workers): - # Create shutdown event for this worker - shutdown_event = Event() - self.shutdown_events.append(shutdown_event) - - # Create process - pass only picklable arguments - process = Process( - target=worker_process_target, - args=(i, self.model_registry, self.process_lock, shutdown_event), - ) - process.daemon = False - process.start() - - self.processes.append(process) - worker_ids.append(i) - - self.logger.debug(f"Started Worker {i} with PID {process.pid}") - - # Start parent monitoring thread - self.monitor_thread = threading.Thread(target=self._parent_monitor, daemon=True) - self.monitor_thread.start() - - msg = f"Started {self.num_workers} worker processes successfully" - self.logger.info(msg) - print(msg) - - return worker_ids - - def shutdown(self, timeout: float = 10.0): - """Shutdown all workers gracefully, then force kill if needed""" - self.logger.debug("WorkerManager shutdown initiated...") - self.shutdown_event.set() - - # Step 1: Signal all workers to shutdown gracefully - self.logger.debug("Signaling graceful shutdown...") - for shutdown_event in self.shutdown_events: - shutdown_event.set() - - # Wait for processes to finish gracefully - start_time = time.time() - all_finished = True - - for i, process in enumerate(self.processes): - if process: - remaining_timeout = max(0, timeout - (time.time() - start_time)) - if remaining_timeout > 0: - process.join(timeout=remaining_timeout) - - if process.is_alive(): - self.logger.debug(f"Worker {i} did not shutdown gracefully") - all_finished = False - - if all_finished: - self.logger.debug("All workers shutdown gracefully") - return - - # Step 2: Send SIGTERM to remaining processes - self.logger.debug("Sending SIGTERM to remaining workers...") - for i, process in enumerate(self.processes): - if process and process.is_alive(): - try: - process.terminate() # Sends SIGTERM - except Exception as e: - self.logger.debug(f"Error sending SIGTERM to Worker {i}: {e}") - - # Wait a bit more - time.sleep(2) - - # Step 3: Force kill any remaining processes - remaining_alive = False - for i, process in enumerate(self.processes): - if process and process.is_alive(): - remaining_alive = True - try: - process.kill() # SIGKILL - process.join(timeout=1.0) - except Exception as e: - self.logger.debug(f"Error force killing Worker {i}: {e}") - - if remaining_alive: - # Nuclear option - kill entire process group - self.logger.debug("Using process group kill as final resort...") - if self.process_group_id: - try: - os.killpg(self.process_group_id, signal.SIGKILL) - except ProcessLookupError: - pass # Process group already dead - except Exception as e: - self.logger.error(f"Error killing process group: {e}") - - self.logger.debug("WorkerManager shutdown complete") - - def get_worker_status(self) -> dict: - """Get status of all workers""" - status = {} - for i, process in enumerate(self.processes): - status[i] = { - "alive": process.is_alive() if process else False, - "pid": process.pid if process else None, - } - return status - - def __del__(self): - """Cleanup when WorkerManager is destroyed""" - try: - self.shutdown(timeout=5.0) - except: - pass diff --git a/rapidfireai/metrics/__init__.py b/rapidfireai/metrics/__init__.py new file mode 100644 index 00000000..15205662 --- /dev/null +++ b/rapidfireai/metrics/__init__.py @@ -0,0 +1,35 @@ +""" +RapidFire AI Metrics Module + +Provides unified metric logging interfaces for different backends +(MLflow, TensorBoard, TrackIO). +""" + +from rapidfireai.metrics.metric_logger import MetricLogger, MetricLoggerConfig, MetricLoggerType +from rapidfireai.metrics.metric_rfmetric_manager import RFMetricLogger + +# Optional imports - these may fail if dependencies aren't installed +try: + from rapidfireai.metrics.metric_mlflow_manager import MLflowMetricLogger +except ImportError: + MLflowMetricLogger = None + +try: + from rapidfireai.metrics.metric_tensorboard_manager import TensorBoardMetricLogger +except ImportError: + TensorBoardMetricLogger = None + +try: + from rapidfireai.metrics.metric_trackio_manager import TrackIOMetricLogger +except ImportError: + TrackIOMetricLogger = None + +__all__ = [ + "MetricLogger", + "MetricLoggerConfig", + "MetricLoggerType", + "MLflowMetricLogger", + "TensorBoardMetricLogger", + "TrackIOMetricLogger", + "RFMetricLogger", +] diff --git a/rapidfireai/utils/metric_logger.py b/rapidfireai/metrics/metric_logger.py similarity index 100% rename from rapidfireai/utils/metric_logger.py rename to rapidfireai/metrics/metric_logger.py diff --git a/rapidfireai/utils/metric_mlflow_manager.py b/rapidfireai/metrics/metric_mlflow_manager.py similarity index 92% rename from rapidfireai/utils/metric_mlflow_manager.py rename to rapidfireai/metrics/metric_mlflow_manager.py index b23020fe..8fecec16 100644 --- a/rapidfireai/utils/metric_mlflow_manager.py +++ b/rapidfireai/metrics/metric_mlflow_manager.py @@ -4,10 +4,10 @@ import mlflow from mlflow.tracking import MlflowClient from typing import Any -from rapidfireai.utils.metric_logger import MetricLogger, MetricLoggerType -from rapidfireai.utils.ping import ping_server +from rapidfireai.metrics.metric_logger import MetricLogger, MetricLoggerType +from rapidfireai.platform.ping import ping_server from rapidfireai.utils.constants import MLflowConfig -from rapidfireai.evals.utils.logger import RFLogger +from rapidfireai.utils.logging import RFLogger class MLflowMetricLogger(MetricLogger): @@ -31,10 +31,11 @@ def __init__(self, tracking_uri: str, logger: RFLogger = None, init_kwargs: dict self.experiment_id = None def create_experiment(self, experiment_name: str) -> str: - """Create a new experiment and set it as active.""" - self.experiment_id = self.client.create_experiment(experiment_name) + """Create a new experiment or get existing one, and set it as active.""" # IMPORTANT: Set this as the active experiment in MLflow context - mlflow.set_experiment(experiment_name) + # mlflow.set_experiment automatically creates if it doesn't exist + experiment = mlflow.set_experiment(experiment_name) + self.experiment_id = experiment.experiment_id return self.experiment_id def get_experiment(self, experiment_name: str) -> str: diff --git a/rapidfireai/utils/metric_rfmetric_manager.py b/rapidfireai/metrics/metric_rfmetric_manager.py similarity index 85% rename from rapidfireai/utils/metric_rfmetric_manager.py rename to rapidfireai/metrics/metric_rfmetric_manager.py index 5fa4559e..a44cd073 100644 --- a/rapidfireai/utils/metric_rfmetric_manager.py +++ b/rapidfireai/metrics/metric_rfmetric_manager.py @@ -4,11 +4,8 @@ from typing import Optional from pathlib import Path -from rapidfireai.utils.metric_logger import MetricLogger, MetricLoggerConfig, MetricLoggerType -from rapidfireai.utils.metric_mlflow_manager import MLflowMetricLogger -from rapidfireai.utils.metric_tensorboard_manager import TensorBoardMetricLogger -from rapidfireai.utils.metric_trackio_manager import TrackioMetricLogger -from rapidfireai.evals.utils.logger import RFLogger +from rapidfireai.metrics.metric_logger import MetricLogger, MetricLoggerConfig, MetricLoggerType +from rapidfireai.utils.logging import RFLogger from rapidfireai.utils.constants import ( MLflowConfig, RF_MLFLOW_ENABLED, @@ -17,6 +14,26 @@ RF_TENSORBOARD_LOG_DIR ) +# Optional imports - these may fail if dependencies aren't installed +MLflowMetricLogger = None +TensorBoardMetricLogger = None +TrackIOMetricLogger = None + +try: + from rapidfireai.metrics.metric_mlflow_manager import MLflowMetricLogger +except ImportError: + pass + +try: + from rapidfireai.metrics.metric_tensorboard_manager import TensorBoardMetricLogger +except ImportError: + pass + +try: + from rapidfireai.metrics.metric_trackio_manager import TrackIOMetricLogger +except ImportError: + pass + class RFMetricLogger(MetricLogger): """ Implementation of MetricLogger that logs to multiple backends. Currently no @@ -39,7 +56,7 @@ def __init__(self, metric_loggers: dict[str, MetricLoggerConfig], logger: RFLogg - "name": {"type": MetricLoggerType.TRACKIO, "config": {"tracking_uri": None}} """ self.type = MetricLoggerType.MULTIPLE - self.logger = logger if logger is not None else RFLogger() + self.logger = logger if logger is not None else RFLogger() if not isinstance(metric_loggers, dict): raise ValueError("metric_loggers must be a dictionary") if len(metric_loggers) == 0: @@ -50,16 +67,25 @@ def __init__(self, metric_loggers: dict[str, MetricLoggerConfig], logger: RFLogg if metric_logger_config.get("type") not in MetricLoggerType: raise ValueError(f"metric_logger_config for {metric_logger_name} must be a valid MetricLoggerType") if metric_logger_config.get("type") == MetricLoggerType.MLFLOW: + if MLflowMetricLogger is None: + self.logger.warning(f"MLflow not available, skipping {metric_logger_name}") + continue try: self.metric_loggers[metric_logger_name] = MLflowMetricLogger(metric_logger_config["config"]["tracking_uri"], logger=self.logger) self.logger.info(f"Initialized MLflowMetricLogger: {metric_logger_name}") except ConnectionRefusedError as e: self.logger.warning(f"Failed to initialize MLflowMetricLogger: {e}. MLflow logging is disabled.") elif metric_logger_config.get("type") == MetricLoggerType.TENSORBOARD: + if TensorBoardMetricLogger is None: + self.logger.warning(f"TensorBoard not available, skipping {metric_logger_name}") + continue self.metric_loggers[metric_logger_name] = TensorBoardMetricLogger(metric_logger_config["config"]["log_dir"], logger=self.logger) self.logger.info(f"Initialized TensorBoardMetricLogger: {metric_logger_name}") elif metric_logger_config.get("type") == MetricLoggerType.TRACKIO: - self.metric_loggers[metric_logger_name] = TrackioMetricLogger( + if TrackIOMetricLogger is None: + self.logger.warning(f"TrackIO not available, skipping {metric_logger_name}") + continue + self.metric_loggers[metric_logger_name] = TrackIOMetricLogger( experiment_name=metric_logger_config["config"]["experiment_name"], logger=self.logger, init_kwargs=metric_logger_config["config"].get("init_kwargs") @@ -67,12 +93,11 @@ def __init__(self, metric_loggers: dict[str, MetricLoggerConfig], logger: RFLogg self.logger.info(f"Initialized TrackioMetricLogger: {metric_logger_name}") else: raise ValueError(f"metric_logger_config for {metric_logger_name} must be a valid MetricLoggerType") - + def _translate_run_name(self, run_id: str) -> str: if len(run_id) == 32: - # Run is a mlflow run id, so we need to get the run name from the mlflow run id try: - from mlflow.tracking import MlflowClient # type: ignore[import-not-found] + from mlflow.tracking import MlflowClient client = MlflowClient(tracking_uri=MLflowConfig.URL) run = client.get_run(run_id) return run.info.run_name @@ -80,15 +105,9 @@ def _translate_run_name(self, run_id: str) -> str: self.logger.warning(f"Error getting run name from mlflow run id {run_id}: {e}, using run id as run name") return run_id return run_id - - def _get_run_name(self, run_id: str) -> str: - """ - Return the user-friendly run name for a canonical run_id, caching the result. - Important: do NOT use dict.setdefault(run_id, self._translate_run_name(run_id)) here, - because the default argument is evaluated eagerly and would trigger MLflow calls on - every invocation. - """ + def _get_run_name(self, run_id: str) -> str: + """Return the user-friendly run name for a canonical run_id, caching the result.""" cached = self.run_id_map.get(run_id) if cached is not None: return cached @@ -101,11 +120,17 @@ def add_logger(self, metric_logger_name: str, metric_logger_config: MetricLogger if metric_logger_config.get("type") not in MetricLoggerType: raise ValueError(f"metric_logger_config for {metric_logger_name} must be a valid MetricLoggerType") if metric_logger_config.get("type") == MetricLoggerType.MLFLOW: + if MLflowMetricLogger is None: + raise ImportError("MLflow is not installed. Install with: pip install mlflow") self.metric_loggers[metric_logger_name] = MLflowMetricLogger(metric_logger_config["config"]["tracking_uri"]) elif metric_logger_config.get("type") == MetricLoggerType.TENSORBOARD: + if TensorBoardMetricLogger is None: + raise ImportError("TensorBoard is not installed. Install with: pip install tensorboard") self.metric_loggers[metric_logger_name] = TensorBoardMetricLogger(metric_logger_config["config"]["log_dir"]) elif metric_logger_config.get("type") == MetricLoggerType.TRACKIO: - self.metric_loggers[metric_logger_name] = TrackioMetricLogger( + if TrackIOMetricLogger is None: + raise ImportError("TrackIO is not installed. Install with: pip install trackio") + self.metric_loggers[metric_logger_name] = TrackIOMetricLogger( experiment_name=metric_logger_config["config"]["experiment_name"], init_kwargs=metric_logger_config["config"].get("init_kwargs") ) @@ -119,14 +144,14 @@ def create_experiment(self, experiment_name: str) -> str: self.logger.info(f"Creating MLflow experiment: {experiment_name}") return metric_logger.create_experiment(experiment_name) return experiment_name - + def get_experiment(self, experiment_name: str) -> str: """Get experiment from MetricLogger(TensorBoard doesn't have experiments).""" for metric_logger in self.metric_loggers.values(): if metric_logger.type == MetricLoggerType.MLFLOW: return metric_logger.get_experiment(experiment_name) return experiment_name - + def create_run(self, run_name: str) -> str: """Create run in MetricLogger. @@ -151,7 +176,7 @@ def create_run(self, run_name: str) -> str: metric_logger.create_run(run_name) return canonical_run_id - + def log_param(self, run_id: str, key: str, value: str) -> None: """Log parameter to MetricLogger.""" run_name = self._get_run_name(run_id) @@ -163,7 +188,7 @@ def log_param(self, run_id: str, key: str, value: str) -> None: metric_logger.log_param(run_name, key, value) else: raise ValueError(f"metric_logger for {metric_logger_name} does not support log_param") - + def log_metric(self, run_id: str, key: str, value: float, step: Optional[int] = None) -> None: """Log metric to MetricLogger.""" run_name = self._get_run_name(run_id) @@ -175,7 +200,7 @@ def log_metric(self, run_id: str, key: str, value: float, step: Optional[int] = metric_logger.log_metric(run_name, key, value, step=step) else: raise ValueError(f"metric_logger for {metric_logger_name} does not support log_metric") - + def get_run_metrics(self, run_id: str) -> dict: """Get metrics from MetricLogger.""" for metric_logger in self.metric_loggers.values(): @@ -208,16 +233,16 @@ def delete_run(self, run_id: str) -> None: else: raise ValueError(f"metric_logger for {metric_logger_name} does not support delete_run") return None - + def clear_context(self) -> None: """Clear context in MetricLogger.""" for metric_logger_name, metric_logger in self.metric_loggers.items(): if hasattr(metric_logger, "clear_context"): metric_logger.clear_context() else: - raise ValueError(f"metric_logger for {metric_logger_name} does not support clear_context") + raise ValueError(f"metric_logger for {metric_logger_name} does not support clear_context") return None - + @classmethod def get_default_metric_loggers(cls, experiment_name: str) -> dict[str, MetricLoggerConfig]: """Get default metric loggers.""" diff --git a/rapidfireai/utils/metric_tensorboard_manager.py b/rapidfireai/metrics/metric_tensorboard_manager.py similarity index 97% rename from rapidfireai/utils/metric_tensorboard_manager.py rename to rapidfireai/metrics/metric_tensorboard_manager.py index e95b9f10..e6774f9a 100644 --- a/rapidfireai/utils/metric_tensorboard_manager.py +++ b/rapidfireai/metrics/metric_tensorboard_manager.py @@ -4,12 +4,12 @@ Uses torch.utils.tensorboard.SummaryWriter to log metrics to TensorBoard. """ -from rapidfireai.utils.metric_logger import MetricLogger, MetricLoggerType +from rapidfireai.metrics.metric_logger import MetricLogger, MetricLoggerType from pathlib import Path from typing import Optional, Any import os from rapidfireai.utils.os_utils import mkdir_p -from rapidfireai.evals.utils.logger import RFLogger +from rapidfireai.utils.logging import RFLogger class TensorBoardMetricLogger(MetricLogger): """ @@ -113,7 +113,7 @@ def get_run_metrics(self, run_id: str) -> dict: This returns an empty dict. For viewing metrics, use TensorBoard UI. """ return {} - + def end_run(self, run_id: str) -> None: """End a TensorBoard run by closing the writer.""" if run_id in self.writers: @@ -156,7 +156,7 @@ def delete_run(self, run_id: str) -> None: destination = os.path.join(deleted_dir, f"{run_id}_{timestamp}") shutil.move(run_log_dir, destination) - + def __del__(self): """Clean up all writers on deletion.""" for writer in self.writers.values(): diff --git a/rapidfireai/utils/metric_trackio_manager.py b/rapidfireai/metrics/metric_trackio_manager.py similarity index 97% rename from rapidfireai/utils/metric_trackio_manager.py rename to rapidfireai/metrics/metric_trackio_manager.py index 176ef767..492d564a 100644 --- a/rapidfireai/utils/metric_trackio_manager.py +++ b/rapidfireai/metrics/metric_trackio_manager.py @@ -5,8 +5,8 @@ import io from contextlib import redirect_stdout from typing import Any -from rapidfireai.utils.metric_logger import MetricLogger, MetricLoggerType -from rapidfireai.evals.utils.logger import RFLogger +from rapidfireai.metrics.metric_logger import MetricLogger, MetricLoggerType +from rapidfireai.utils.logging import RFLogger import warnings warnings.filterwarnings("ignore", message="Reserved keys renamed") @@ -30,6 +30,7 @@ def __init__(self, experiment_name: str, logger: RFLogger = None, init_kwargs: d self.logger = logger if logger is not None else RFLogger() self.active_runs = {} # Map run_id -> runs self.run_params = {} # Map run_id -> dict of params to log on init + self._initialized = False def _capture_trackio_output(self, func, *args, **kwargs): """ @@ -89,7 +90,7 @@ def create_run(self, run_name: str) -> str: f"with self.init_kwargs={self.init_kwargs!r}: {exc}" ) from exc - # Log any pending params for this run + # Log any pending params for this run if run_name in self.run_params: self.active_runs[run_name].log(self.run_params[run_name]) del self.run_params[run_name] @@ -130,7 +131,7 @@ def log_metric(self, run_id: str, key: str, value: float, step: int = None) -> N def get_run_metrics(self, run_id: str) -> dict[str, list[tuple[int, float]]]: """ Get all metrics for a specific run. - + Note: Trackio stores metrics locally. This method returns an empty dict as Trackio doesn't provide a direct API to retrieve historical metrics. Metrics can be viewed using `trackio.show()`. diff --git a/rapidfireai/platform/__init__.py b/rapidfireai/platform/__init__.py new file mode 100644 index 00000000..eaee5525 --- /dev/null +++ b/rapidfireai/platform/__init__.py @@ -0,0 +1,25 @@ +""" +RapidFire AI Platform Module + +Provides platform/environment detection, diagnostics, and health check utilities. +""" + +from rapidfireai.platform.colab import is_running_in_colab, get_colab_auth_token +from rapidfireai.platform.doctor import get_doctor_info +from rapidfireai.platform.get_ip_address import get_ip_address +from rapidfireai.platform.gpu_info import get_gpu_info, get_compute_capability, get_torch_version +from rapidfireai.platform.python_info import get_python_info, get_pip_packages +from rapidfireai.platform.ping import ping_server + +__all__ = [ + "is_running_in_colab", + "get_colab_auth_token", + "get_doctor_info", + "get_ip_address", + "get_gpu_info", + "get_compute_capability", + "get_torch_version", + "get_python_info", + "get_pip_packages", + "ping_server", +] diff --git a/rapidfireai/utils/colab.py b/rapidfireai/platform/colab.py similarity index 100% rename from rapidfireai/utils/colab.py rename to rapidfireai/platform/colab.py diff --git a/rapidfireai/utils/doctor.py b/rapidfireai/platform/doctor.py similarity index 94% rename from rapidfireai/utils/doctor.py rename to rapidfireai/platform/doctor.py index a676cebc..5768d284 100644 --- a/rapidfireai/utils/doctor.py +++ b/rapidfireai/platform/doctor.py @@ -5,27 +5,29 @@ import os import platform from pathlib import Path -from rapidfireai.utils.python_info import get_python_info, get_pip_packages -from rapidfireai.utils.gpu_info import get_gpu_info, get_torch_version -from rapidfireai.utils.ping import ping_server + from rapidfireai.utils.constants import ( - JupyterConfig, + ColabConfig, DispatcherConfig, - MLflowConfig, FrontendConfig, + JupyterConfig, + MLflowConfig, RayConfig, - ColabConfig, + RF_DB_PATH, + RF_EXPERIMENT_PATH, RF_HOME, + RF_LOG_FILENAME, RF_LOG_PATH, - RF_EXPERIMENT_PATH, - RF_DB_PATH, - RF_TENSORBOARD_LOG_DIR, - RF_TRAINING_LOG_FILENAME, RF_MLFLOW_ENABLED, RF_TENSORBOARD_ENABLED, + RF_TENSORBOARD_LOG_DIR, RF_TRACKIO_ENABLED, - RF_LOG_FILENAME, + RF_TRAINING_LOG_FILENAME, ) +from rapidfireai.platform.gpu_info import get_gpu_info, get_torch_version +from rapidfireai.platform.ping import ping_server +from rapidfireai.platform.python_info import get_pip_packages, get_python_info + def get_doctor_info(log_lines: int = 10): """ @@ -157,7 +159,7 @@ def get_doctor_info(log_lines: int = 10): print(f"Torch Version: {major}.{minor}.{patch}") else: status = 1 if status == 0 else status - print("โš ๏ธ Torch version not found") + print("โš ๏ธ Torch version not found") if int(torch_cuda_major) > 0: print(f"Torch CUDA Version: {torch_cuda_major}.{torch_cuda_minor}") else: @@ -165,15 +167,15 @@ def get_doctor_info(log_lines: int = 10): print("โš ๏ธ Torch CUDA Version: unknown") # Check RapidFire AI ports - print ("\n๐Ÿ›œ RapidFire AI Ports:") - print ("-" * 30) + print("\n๐Ÿ›œ RapidFire AI Ports:") + print("-" * 30) for check_item in [ - {"host": JupyterConfig.HOST, "port": JupyterConfig.PORT, "service": "Jupyter"}, - {"host": DispatcherConfig.HOST, "port": DispatcherConfig.PORT, "service": "API(Dispatcher)"}, + {"host": JupyterConfig.HOST, "port": JupyterConfig.PORT, "service": "Jupyter"}, + {"host": DispatcherConfig.HOST, "port": DispatcherConfig.PORT, "service": "API(Dispatcher)"}, {"host": MLflowConfig.HOST, "port": MLflowConfig.PORT, "service": "MLflow"}, {"host": FrontendConfig.HOST, "port": FrontendConfig.PORT, "service": "Frontend"}, - {"host": RayConfig.HOST, "port": RayConfig.PORT, "service": "Ray"}]: - + {"host": RayConfig.HOST, "port": RayConfig.PORT, "service": "Ray"}, + ]: for host_index, host_check in enumerate(["127.0.0.1", check_item["host"]]): if host_index == 0 or (host_check not in ["127.0.0.1", "0.0.0.0"]): checker = ping_server(host_check, check_item["port"]) @@ -181,7 +183,7 @@ def get_doctor_info(log_lines: int = 10): print(f"{check_item['service']}: is currently Listening on {host_check}:{check_item['port']}") else: print(f"{check_item['service']}: is NOT currently listening on {host_check}:{check_item['port']}") - + # System Information print("\n๐Ÿ’ป System Information:") print("-" * 30) @@ -227,7 +229,7 @@ def get_doctor_info(log_lines: int = 10): print(f"\n๐Ÿ“„{current_item}:") if log_lines != 0: if not os.path.isdir(current_item): - with open(current_item, "r", encoding="utf-8") as f: + with open(current_item, encoding="utf-8") as f: lines = f.readlines() read_lines = lines[-log_lines:] if log_lines == -1: @@ -242,4 +244,4 @@ def get_doctor_info(log_lines: int = 10): elif status == 2: print("\nโŒ Diagnostics complete with errors") else: - print("\nโŒ Diagnostics completed with unknown status") \ No newline at end of file + print("\nโŒ Diagnostics completed with unknown status") diff --git a/rapidfireai/utils/get_ip_address.py b/rapidfireai/platform/get_ip_address.py similarity index 100% rename from rapidfireai/utils/get_ip_address.py rename to rapidfireai/platform/get_ip_address.py diff --git a/rapidfireai/utils/gpu_info.py b/rapidfireai/platform/gpu_info.py similarity index 97% rename from rapidfireai/utils/gpu_info.py rename to rapidfireai/platform/gpu_info.py index 1bdd6548..7276a98f 100644 --- a/rapidfireai/utils/gpu_info.py +++ b/rapidfireai/platform/gpu_info.py @@ -1,9 +1,10 @@ """Utility functions for GPU information.""" +import os +import re import shutil import subprocess -import re -import os + def get_compute_capability(): """Get compute capability from nvidia-smi""" @@ -22,6 +23,7 @@ def get_compute_capability(): except (subprocess.CalledProcessError, FileNotFoundError): return None + def get_gpu_info(): """Get comprehensive GPU and CUDA information.""" info = {"status": 0} @@ -137,10 +139,13 @@ def get_gpu_info(): return info + def get_torch_version(): """Get torch major, minor, patch version, along with cuda version if installed""" try: - result = subprocess.run(["python", "-c", "import torch; print(torch.__version__)"], capture_output=True, text=True, check=True) + result = subprocess.run( + ["python", "-c", "import torch; print(torch.__version__)"], capture_output=True, text=True, check=True + ) version = result.stdout.strip() # version maybe like 2.8.0+cu128 or 2.8.0 cuda_major = "0" @@ -152,4 +157,4 @@ def get_torch_version(): major, minor, patch = version.split("+")[0].split(".") return major, minor, patch, cuda_major, cuda_minor except (subprocess.CalledProcessError, FileNotFoundError, IndexError): - return "0","0","0","0","0" \ No newline at end of file + return "0", "0", "0", "0", "0" diff --git a/rapidfireai/utils/ping.py b/rapidfireai/platform/ping.py similarity index 100% rename from rapidfireai/utils/ping.py rename to rapidfireai/platform/ping.py diff --git a/rapidfireai/utils/python_info.py b/rapidfireai/platform/python_info.py similarity index 95% rename from rapidfireai/utils/python_info.py rename to rapidfireai/platform/python_info.py index 13bc2579..0320bf5f 100644 --- a/rapidfireai/utils/python_info.py +++ b/rapidfireai/platform/python_info.py @@ -1,10 +1,11 @@ """Utility functions for Python information.""" import os -import sys import platform import site import subprocess +import sys + def get_python_info(): """Get comprehensive Python information.""" @@ -26,10 +27,11 @@ def get_python_info(): return info + def get_pip_packages(): """Get list of installed pip packages.""" try: result = subprocess.run([sys.executable, "-m", "pip", "list"], capture_output=True, text=True, check=True) return result.stdout except (subprocess.CalledProcessError, FileNotFoundError): - return "Failed to get pip packages" \ No newline at end of file + return "Failed to get pip packages" diff --git a/rapidfireai/utils/__init__.py b/rapidfireai/utils/__init__.py index fb5ef39e..d7fdb8ad 100644 --- a/rapidfireai/utils/__init__.py +++ b/rapidfireai/utils/__init__.py @@ -1,7 +1,110 @@ -"""Utility functions and helpers.""" +""" +RapidFire AI Utility functions and helpers. -from .colab import get_colab_auth_token +This module provides unified utilities for both fit and evals modes. +""" + +from rapidfireai.platform.colab import get_colab_auth_token, is_running_in_colab +from .constants import ( + # Configs + DispatcherConfig, + FrontendConfig, + MLflowConfig, + JupyterConfig, + RayConfig, + ColabConfig, + # Paths + RF_HOME, + RF_LOG_PATH, + RF_LOG_FILENAME, + RF_EXPERIMENT_PATH, + RF_DB_PATH, + RF_TRAINING_LOG_FILENAME, + # Enums + ExperimentStatus, + RunStatus, + PipelineStatus, + ContextStatus, + TaskStatus, + ICOperation, + ICStatus, + ExperimentTask, + ControllerTask, + WorkerTask, + RunSource, + RunEndedBy, + LogType, + SHMObjectType, + # Functions + get_dispatcher_url, + get_dispatcher_headers, +) +from .exceptions import ( + RFException, + ExperimentException, + DispatcherException, + DBException, + ControllerException, + WorkerException, + PipelineException, + ActorException, +) +from .logging import RFLogger, TrainingLogger +from .serialize import encode_payload, decode_db_payload, extract_pipeline_config_json +from .experiment_utils import ExperimentUtils __all__ = [ + # Colab "get_colab_auth_token", + "is_running_in_colab", + # Configs + "DispatcherConfig", + "FrontendConfig", + "MLflowConfig", + "JupyterConfig", + "RayConfig", + "ColabConfig", + # Paths + "RF_HOME", + "RF_LOG_PATH", + "RF_LOG_FILENAME", + "RF_EXPERIMENT_PATH", + "RF_DB_PATH", + "RF_TRAINING_LOG_FILENAME", + # Enums + "ExperimentStatus", + "RunStatus", + "PipelineStatus", + "ContextStatus", + "TaskStatus", + "ICOperation", + "ICStatus", + "ExperimentTask", + "ControllerTask", + "WorkerTask", + "RunSource", + "RunEndedBy", + "LogType", + "SHMObjectType", + # Functions + "get_dispatcher_url", + "get_dispatcher_headers", + # Exceptions + "RFException", + "ExperimentException", + "DispatcherException", + "DBException", + "ControllerException", + "WorkerException", + "PipelineException", + "ActorException", + # Logging + "RFLogger", + "TrainingLogger", + # Serialization + "encode_payload", + "decode_db_payload", + "extract_pipeline_config_json", + # Experiment + "ExperimentUtils", ] diff --git a/rapidfireai/utils/constants.py b/rapidfireai/utils/constants.py index d1df7fc4..5eac439f 100644 --- a/rapidfireai/utils/constants.py +++ b/rapidfireai/utils/constants.py @@ -1,19 +1,14 @@ """ Constants for the RapidFire AI package + +This module contains constants, configuration classes, and enums. """ + import os from enum import Enum -from rapidfireai.utils.colab import is_running_in_colab -from rapidfireai.utils.os_utils import mkdir_p - - -class ExperimentStatus(str, Enum): - """Shared status values for experiments (used by both fit and evals).""" - RUNNING = "running" - COMPLETED = "completed" - FAILED = "failed" - CANCELLED = "cancelled" +from rapidfireai.platform.colab import is_running_in_colab +from rapidfireai.utils.os_utils import mkdir_p if is_running_in_colab(): @@ -28,6 +23,47 @@ class ExperimentStatus(str, Enum): RF_TRAINING_LOG_FILENAME = os.getenv("RF_TRAINING_LOG_FILENAME", "training.log") RF_TRAINER_OUTPUT = os.getenv("RF_TRAINER_OUTPUT", os.path.join(RF_HOME, "trainer_output")) +# Evals Actor Constants +NUM_QUERY_PROCESSING_ACTORS = 4 +NUM_CPUS_PER_DOC_ACTOR = 2 if os.cpu_count() > 2 else 1 + +# Rate Limiting Constants +MAX_RATE_LIMIT_RETRIES = 5 +RATE_LIMIT_BACKOFF_BASE = 2 + +# --------------------------------------------------------------------------- +# Reranker class registry +# --------------------------------------------------------------------------- +def _build_reranker_registry() -> dict[str, type]: + registry: dict[str, type] = {} + try: + from langchain_classic.retrievers.document_compressors import CrossEncoderReranker + registry[CrossEncoderReranker.__qualname__] = CrossEncoderReranker + except ImportError: + pass + try: + from langchain.retrievers.document_compressors import LLMChainExtractor + registry[LLMChainExtractor.__qualname__] = LLMChainExtractor + except ImportError: + pass + try: + from langchain_cohere import CohereRerank + registry[CohereRerank.__qualname__] = CohereRerank + except ImportError: + pass + return registry + +RERANKER_CLASS_REGISTRY: dict[str, type] = _build_reranker_registry() + +# RAG search type defaults +VALID_SEARCH_TYPES = {"similarity", "similarity_score_threshold", "mmr"} +SEARCH_DEFAULTS: dict[str, dict] = { + "similarity": {"k": 5, "filter": None}, + "similarity_score_threshold": {"k": 5, "filter": None, "score_threshold": 0.5}, + "mmr": {"k": 5, "filter": None, "fetch_k": 20, "lambda_mult": 0.5}, +} +SEARCH_TYPE_KEYS: dict[str, set] = {search_type: set(defaults.keys()) for search_type, defaults in SEARCH_DEFAULTS.items()} + try: mkdir_p(RF_LOG_PATH) mkdir_p(RF_HOME) @@ -35,6 +71,7 @@ class ExperimentStatus(str, Enum): print(f"Error creating directory: {e}") raise + class DispatcherConfig: """Class to manage the dispatcher configuration""" @@ -45,6 +82,7 @@ class DispatcherConfig: def __str__(self): return f"DispatcherConfig(HOST={self.HOST}, PORT={self.PORT}, URL={self.URL})" + # Frontend Constants class FrontendConfig: """Class to manage the frontend configuration""" @@ -56,6 +94,7 @@ class FrontendConfig: def __str__(self): return f"FrontendConfig(HOST={self.HOST}, PORT={self.PORT}, URL={self.URL})" + # MLflow Constants class MLflowConfig: """Class to manage the MLflow configuration""" @@ -67,6 +106,7 @@ class MLflowConfig: def __str__(self): return f"MLflowConfig(HOST={self.HOST}, PORT={self.PORT}, URL={self.URL})" + # Jupyter Constants class JupyterConfig: """Class to manage the Jupyter configuration""" @@ -78,6 +118,7 @@ class JupyterConfig: def __str__(self): return f"JupyterConfig(HOST={self.HOST}, PORT={self.PORT}, URL={self.URL})" + # Ray Constants class RayConfig: """Class to manage the Ray configuration""" @@ -89,6 +130,7 @@ class RayConfig: def __str__(self): return f"RayConfig(HOST={self.HOST}, PORT={self.PORT}, URL={self.URL})" + # Colab Constants class ColabConfig: """Class to manage the Colab configuration""" @@ -99,6 +141,198 @@ class ColabConfig: def __str__(self): return f"ColabConfig(ON_COLAB={self.ON_COLAB}, RF_COLAB_MODE={self.RF_COLAB_MODE})" + RF_MLFLOW_ENABLED = os.getenv("RF_MLFLOW_ENABLED", "true" if not ColabConfig.ON_COLAB else "false") RF_TENSORBOARD_ENABLED = os.getenv("RF_TENSORBOARD_ENABLED", "false" if not ColabConfig.ON_COLAB else "true") -RF_TRACKIO_ENABLED = os.getenv("RF_TRACKIO_ENABLED", "false") \ No newline at end of file +RF_TRACKIO_ENABLED = os.getenv("RF_TRACKIO_ENABLED", "false") +RF_TRACKING_BACKEND = os.getenv("RF_TRACKING_BACKEND", "mlflow" if not ColabConfig.ON_COLAB else "tensorboard") + + +# ============================================================================ +# UNIFIED STATUS ENUMS (lowercase values for consistency) +# ============================================================================ + + +class ExperimentStatus(str, Enum): + """Status for experiments (both fit and evals).""" + + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + + +class RunStatus(str, Enum): + """Status for fit training runs.""" + + NEW = "new" + ONGOING = "ongoing" + COMPLETED = "completed" + STOPPED = "stopped" + DELETED = "deleted" + FAILED = "failed" + + +class PipelineStatus(str, Enum): + """Status for evals inference pipelines.""" + + NEW = "new" + ONGOING = "ongoing" + COMPLETED = "completed" + STOPPED = "stopped" + DELETED = "deleted" + FAILED = "failed" + + +class ContextStatus(str, Enum): + """Status for RAG contexts (evals mode).""" + + NEW = "new" + ONGOING = "ongoing" + COMPLETED = "completed" + DELETED = "deleted" + FAILED = "failed" + + +class TaskStatus(str, Enum): + """Status for worker tasks (fit) and actor tasks (evals).""" + + SCHEDULED = "scheduled" + IN_PROGRESS = "in_progress" + COMPLETED = "completed" + FAILED = "failed" + SKIPPED = "skipped" # Fit-specific + + +class ICOperation(str, Enum): + """Interactive control operation types.""" + + STOP = "stop" + RESUME = "resume" + DELETE = "delete" + CLONE = "clone" + CLONE_WARM = "clone_warm" # Fit-specific (warm-start) + + +class ICStatus(str, Enum): + """Status for interactive control operations.""" + + PENDING = "pending" + PROCESSING = "processing" + COMPLETED = "completed" + FAILED = "failed" + SKIPPED = "skipped" + + +# ============================================================================ +# FIT-SPECIFIC ENUMS +# ============================================================================ + + +class ExperimentTask(str, Enum): + """Fit-mode experiment tasks (current_task column).""" + + IDLE = "idle" + CREATE_MODELS = "create_models" + IC_OPS = "ic_ops" + RUN_FIT = "run_fit" + + +class ControllerTask(str, Enum): + """Fit-mode controller tasks.""" + + RUN_FIT = "run_fit" + CREATE_MODELS = "create_models" + IC_DELETE = "ic_delete" + IC_STOP = "ic_stop" + IC_RESUME = "ic_resume" + IC_CLONE_MODIFY = "ic_clone_modify" + IC_CLONE_MODIFY_WARM = "ic_clone_modify_warm" + EPOCH_BOUNDARY = "epoch_boundary" + GET_RUN_METRICS = "get_run_metrics" + + +class WorkerTask(str, Enum): + """Fit-mode worker tasks.""" + + CREATE_MODELS = "create_models" + TRAIN_VAL = "train_val" + + +class RunSource(str, Enum): + """How a fit run was created.""" + + SHA = "sha" # Successive Halving Algorithm + INITIAL = "initial" + INTERACTIVE_CONTROL = "interactive_control" + + +class RunEndedBy(str, Enum): + """How a fit run was ended.""" + + SHA = "sha" # Successive Halving Algorithm + EPOCH_COMPLETED = "epoch_completed" + INTERACTIVE_CONTROL = "interactive_control" + TOLERANCE = "tolerance" + + +class LogType(str, Enum): + """Log file types.""" + + RF_LOG = "rf_log" + TRAINING_LOG = "training_log" + + +class SHMObjectType(str, Enum): + """Types of objects stored in shared memory (fit mode).""" + + BASE_MODEL = "base_model" + FULL_MODEL = "full_model" + REF_FULL_MODEL = "ref_full_model" + REF_STATE_DICT = "ref_state_dict" + CHECKPOINTS = "checkpoints" + + +# ============================================================================ +# HELPER FUNCTIONS +# ============================================================================ + + +def get_dispatcher_url() -> str: + """ + Auto-detect dispatcher URL based on environment. + + Returns: + - In Google Colab: Uses Colab's kernel proxy URL + - In Jupyter/Local: Uses localhost URL + """ + if ColabConfig.ON_COLAB: + try: + from google.colab.output import eval_js + + proxy_url = eval_js(f"google.colab.kernel.proxyPort({DispatcherConfig.PORT})") + print(f"๐ŸŒ Google Colab detected. Dispatcher URL: {proxy_url}") + return proxy_url + except Exception as e: + print(f"โš ๏ธ Colab detected but failed to get proxy URL: {e}") + return DispatcherConfig.URL + else: + return DispatcherConfig.URL + + +def get_dispatcher_headers() -> dict[str, str]: + """ + Get the HTTP headers needed for dispatcher API requests. + + Returns: + Dictionary with required headers, including Authorization header in Colab + """ + from rapidfireai.platform.colab import get_colab_auth_token + + headers = {"Content-Type": "application/json"} + + auth_token = get_colab_auth_token() + if auth_token: + headers["Authorization"] = f"Bearer {auth_token}" + + return headers diff --git a/rapidfireai/utils/dispatcher_utils.py b/rapidfireai/utils/dispatcher_utils.py index 61903c6d..4dfcc3fb 100644 --- a/rapidfireai/utils/dispatcher_utils.py +++ b/rapidfireai/utils/dispatcher_utils.py @@ -4,6 +4,8 @@ from typing import Any, Protocol +from rapidfireai.utils.constants import ExperimentStatus + class DatabaseWithExperiment(Protocol): """Protocol for database classes that support get_running_experiment.""" @@ -26,8 +28,12 @@ def check_experiment_running(db: DatabaseWithExperiment, experiment_name: str) - """ try: running_experiment = db.get_running_experiment() - running_name = running_experiment.get("experiment_name") if running_experiment else None - running_status = running_experiment.get("status") if running_experiment else None - return running_name == experiment_name and running_status == "running" + if not running_experiment: + return False + running_name = running_experiment.get("experiment_name") + running_status = running_experiment.get("status") + # Compare with enum - works for both enum values and strings due to str inheritance + is_running = running_status == ExperimentStatus.RUNNING + return running_name == experiment_name and is_running except Exception: return False diff --git a/rapidfireai/utils/distributed_utils.py b/rapidfireai/utils/distributed_utils.py index fd4a39e8..6e1153cb 100644 --- a/rapidfireai/utils/distributed_utils.py +++ b/rapidfireai/utils/distributed_utils.py @@ -28,25 +28,32 @@ def setup_distributed_environment( master_addr: str = "localhost", master_port: int | None = None, timeout_minutes: int = 30, + local_rank: int | None = None, ) -> None: """ Initialize distributed training environment. + + Args: + rank: Global rank of this process in the distributed group. + world_size: Total number of processes in the group. + local_rank: Local device index. Defaults to rank for multiprocessing, + but should be 0 when each Ray actor has exactly 1 GPU assigned. """ - # Check if already initialized if dist.is_initialized(): return + if local_rank is None: + local_rank = rank + if master_port is None: master_port = find_free_port() - # Set environment variables for distributed training os.environ["MASTER_ADDR"] = master_addr os.environ["MASTER_PORT"] = str(master_port) os.environ["RANK"] = str(rank) os.environ["WORLD_SIZE"] = str(world_size) - os.environ["LOCAL_RANK"] = str(rank) + os.environ["LOCAL_RANK"] = str(local_rank) - # Initialize the process group timeout = torch.distributed.constants.default_pg_timeout if timeout_minutes > 0: timeout = torch.distributed.constants.default_pg_timeout * timeout_minutes @@ -56,11 +63,11 @@ def setup_distributed_environment( rank=rank, world_size=world_size, timeout=timeout, - device_id=torch.device(f"cuda:{rank}") if torch.cuda.is_available() else None, + device_id=torch.device(f"cuda:{local_rank}") if torch.cuda.is_available() else None, ) if torch.cuda.is_available(): - torch.cuda.set_device(rank) + torch.cuda.set_device(local_rank) def cleanup_distributed() -> None: diff --git a/rapidfireai/utils/exceptions.py b/rapidfireai/utils/exceptions.py new file mode 100644 index 00000000..98e4eba0 --- /dev/null +++ b/rapidfireai/utils/exceptions.py @@ -0,0 +1,99 @@ +""" +RapidFire AI Exceptions + +Custom exceptions for RapidFire AI operations. +""" + + +class RFException(Exception): + """Base exception for RapidFire AI.""" + + def __init__(self, message: str): + self.message = message + super().__init__(self.message) + + +class ExperimentException(RFException): + """Exception for experiment creation and management.""" + pass + + +class DispatcherException(RFException): + """Exception for dispatcher operations.""" + pass + + +class DBException(RFException): + """Exception for database operations.""" + pass + + +class DataPathException(RFException): + """Exception for data path operations.""" + pass + + +class NoGPUsFoundException(RFException): + """Exception for no GPUs found.""" + pass + + +class InitializeRunException(RFException): + """Exception for run initialization.""" + pass + + +class ControllerException(RFException): + """Exception for controller operations.""" + pass + + +class WorkerException(RFException): + """Exception for worker operations.""" + + def __init__(self, message: str, traceback: str = None): + self.traceback_str = traceback + super().__init__(message) + + +class AutoMLException(RFException): + """Exception for AutoML operations.""" + pass + + +class InsufficientSharedMemoryException(RFException): + """Exception for insufficient shared memory.""" + pass + + +class PipelineException(RFException): + """Exception for pipeline operations (evals mode).""" + pass + + +class ActorException(RFException): + """Exception for actor operations (evals mode).""" + pass + + +class RateLimitException(RFException): + """Exception for rate limiting errors (evals mode).""" + pass + + +__all__ = [ + "RFException", + "ExperimentException", + "DispatcherException", + "DBException", + "DataPathException", + "NoGPUsFoundException", + "InitializeRunException", + "ControllerException", + "WorkerException", + "AutoMLException", + "InsufficientSharedMemoryException", + "PipelineException", + "ActorException", + "RateLimitException", +] diff --git a/rapidfireai/fit/utils/experiment_utils.py b/rapidfireai/utils/experiment_utils.py similarity index 51% rename from rapidfireai/fit/utils/experiment_utils.py rename to rapidfireai/utils/experiment_utils.py index e7a0401e..0e159344 100644 --- a/rapidfireai/fit/utils/experiment_utils.py +++ b/rapidfireai/utils/experiment_utils.py @@ -1,112 +1,119 @@ -"""This module contains utility functions for the experiment.""" +""" +RapidFire AI Experiment Utilities + +Provides experiment creation and management. +""" -import multiprocessing as mp import os import re -import signal -import sys import warnings -from typing import Any - +from pathlib import Path import pandas as pd -import torch -from IPython.display import display -from tqdm import tqdm -from transformers import logging as transformers_logging - -from rapidfireai.utils.constants import MLflowConfig, RF_MLFLOW_ENABLED -from rapidfireai.utils.metric_rfmetric_manager import RFMetricLogger -from rapidfireai.fit.db.rf_db import RfDb -from rapidfireai.fit.utils.constants import ExperimentStatus, ExperimentTask -from rapidfireai.fit.utils.datapaths import DataPath -from rapidfireai.fit.utils.exceptions import DBException, ExperimentException -from rapidfireai.fit.utils.logging import RFLogger -# Note: mlflow and MLflowManager are imported lazily inside conditional blocks -# to avoid MLflow connection attempts when using tensorboard-only mode +from rapidfireai.db.rf_db import RfDb +from rapidfireai.utils.constants import ( + ExperimentStatus, + ExperimentTask, + RF_EXPERIMENT_PATH, + RF_MLFLOW_ENABLED, +) +from rapidfireai.utils.exceptions import ExperimentException +from rapidfireai.utils.logging import RFLogger +from rapidfireai.utils.os_utils import mkdir_p class ExperimentUtils: - """Class to contain utility functions for the experiment.""" + """ + Experiment utilities for creation, naming, cancellation, and cleanup. + """ - def __init__(self) -> None: - # initialize database handler + def __init__(self): + """Initialize with database handler.""" self.db = RfDb() def _disable_ml_warnings_display(self) -> None: - """Disable notebook display""" - tqdm.disable = True + """Disable ML-related warnings and verbose output.""" + try: + from tqdm import tqdm + tqdm.disable = True + except ImportError: + pass - # Suppress the transformers logging - os.environ["TRANSFORMERS_VERBOSITY"] = "error" - transformers_logging.set_verbosity_error() + try: + import torch + torch.set_warn_always(False) + except ImportError: + pass - # Suppress the torch warnings - torch.set_warn_always(False) + try: + from transformers import logging as transformers_logging + os.environ["TRANSFORMERS_VERBOSITY"] = "error" + transformers_logging.set_verbosity_error() + except ImportError: + pass - # Suppress the FutureWarning + # Suppress common warnings warnings.filterwarnings("ignore", message=".*torch.cuda.amp.autocast.*") warnings.filterwarnings("ignore", message=".*torch.amp.autocast.*") warnings.filterwarnings("ignore", message=".*fan_in_fan_out is set to False.*") 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( + warnings.filterwarnings("ignore", message=".*Unable to register cuDNN factory.*") + warnings.filterwarnings("ignore", message=".*Unable to register cuBLAS factory.*") + warnings.filterwarnings("ignore", message=".*All log messages before absl::InitializeLog.*") + warnings.filterwarnings("ignore", message=".*resource_tracker: process died unexpectedly.*") + warnings.filterwarnings("ignore", message=".*computation placer already registered.*") + warnings.filterwarnings("ignore", message=".*Rank 0 is connected to 0 peer ranks.*") + warnings.filterwarnings("ignore", message=".*No cudagraph will be used.*") + warnings.filterwarnings("ignore", module="multiprocessing.resource_tracker") + + def create_experiment( 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) + given_name: str, + experiments_path: str, + ) -> tuple[int, str, list[str]]: + """ + Create a new experiment. Returns the experiment id, name, and log messages. - print("Graceful shutdown completed.") - sys.exit(0) - except Exception as e: - print(f"Error during graceful shutdown: {e}") - sys.exit(1) + This method handles: + - Checking for already running experiments + - Generating unique experiment names + - Creating experiment record in database + - Creating experiment directories - # Register signal handlers - signal.signal(signal.SIGINT, signal_handler) - signal.signal(signal.SIGTERM, signal_handler) + Args: + given_name: Desired experiment name + experiments_path: Path to experiments directory - 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.""" + Returns: + Tuple of (experiment_id, experiment_name, log_messages) + """ log_messages: list[str] = [] - # disable warnings from notebook + # Disable warnings self._disable_ml_warnings_display() # Clear any existing MLflow context before starting new experiment - # Only if using MLflow backend - if RF_MLFLOW_ENABLED=="true": - import mlflow # Lazy import to avoid connection attempts in tensorboard-only mode - + if RF_MLFLOW_ENABLED == "true": try: + import mlflow 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 + # Check if experiment is already running running_experiment = None try: running_experiment = self.db.get_running_experiment() - except DBException: + except Exception: pass + if running_experiment: - # check if the running experiment is the same as the new experiment if given_name == running_experiment["experiment_name"]: + # Same experiment - reuse it msg = ( f"Experiment {running_experiment['experiment_name']} is currently running." " Returning the same experiment object." @@ -114,20 +121,22 @@ def create_experiment(self, given_name: str, experiments_path: str) -> tuple[int 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) + # Check if there's a running task and cancel it + try: + 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) + except Exception: + pass + self.cancel_current(internal=True) - # get experiment id - experiment_id, experiment_name = ( - running_experiment["experiment_id"], - running_experiment["experiment_name"], - ) + experiment_id = running_experiment["experiment_id"] + experiment_name = running_experiment["experiment_name"] else: + # Different experiment - end the previous one self.end_experiment(internal=True) experiment_id, experiment_name, metric_experiment_id = self._create_experiment_internal( given_name, @@ -143,12 +152,12 @@ def create_experiment(self, given_name: str, experiments_path: str) -> tuple[int msg = ( f"The previously running experiment {running_experiment['experiment_name']} was forcibly ended." f" Created a new experiment '{experiment_name}' with Experiment ID: {experiment_id}" - f" at {experiments_path}/{experiment_name} (TensorBoard-only mode)" + f" 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(): + # Name exists - create with incremented suffix experiment_id, experiment_name, metric_experiment_id = self._create_experiment_internal( given_name, experiments_path, @@ -163,11 +172,12 @@ def create_experiment(self, given_name: str, experiments_path: str) -> tuple[int msg = ( "An experiment with the same name already exists." f" Created a new experiment '{experiment_name}' with Experiment ID: {experiment_id}" - f" at {experiments_path}/{experiment_name} (TensorBoard-only mode)" + f" at {experiments_path}/{experiment_name}" ) print(msg) log_messages.append(msg) else: + # New experiment experiment_id, experiment_name, metric_experiment_id = self._create_experiment_internal( given_name, experiments_path, @@ -180,115 +190,124 @@ def create_experiment(self, given_name: str, experiments_path: str) -> tuple[int else: msg = ( f"Experiment {experiment_name} created with Experiment ID: {experiment_id}" - f" at {experiments_path}/{experiment_name} (TensorBoard-only mode)" + f" at {experiments_path}/{experiment_name}" ) print(msg) log_messages.append(msg) - # initialize the data paths and create directories + # Create experiment directory try: - DataPath.initialize(experiment_name, experiments_path) + experiment_dir = Path(experiments_path) / experiment_name + mkdir_p(experiment_dir) except Exception as e: - raise ExperimentException(f"Failed to initialize data paths: {e}") + raise ExperimentException(f"Failed to create experiment directories: {e}") from e return experiment_id, experiment_name, log_messages def end_experiment(self, internal: bool = False) -> None: - """End the experiment""" - # check if experiment is running + """ + End the experiment and clean up resources. + + Args: + internal: If True, suppress output messages + """ try: current_experiment = self.db.get_running_experiment() - except DBException: + except Exception: 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) + experiments_path = current_experiment.get("experiments_path", RF_EXPERIMENT_PATH) + + # Create logger + logger = RFLogger( + experiment_name=experiment_name, + experiment_path=experiments_path, + ).get_logger("ExperimentUtils") - # cancel current task if there's any + # Cancel current task if any self.cancel_current(internal=True) - # reset DB states + # Reset DB states self.db.set_experiment_status(current_experiment["experiment_id"], ExperimentStatus.COMPLETED) self.db.reset_all_tables() - # Clear MLflow context only if using MLflow backend - if RF_MLFLOW_ENABLED=="true": - import mlflow # Lazy import to avoid connection attempts in tensorboard-only mode - + # Clear MLflow context if enabled + if RF_MLFLOW_ENABLED == "true": try: + import mlflow if mlflow.active_run(): print("Ending active MLflow run before ending experiment") mlflow.end_run() - - # Also clear context through RFMetricLogger if available - try: - metric_logger_config = RFMetricLogger.get_default_metric_loggers(experiment_name=experiment_name) - metric_logger = RFMetricLogger(metric_logger_config, logger=logger) - metric_logger.clear_context() - except Exception as e2: - print(f"Error clearing Metric context through RFMetricLogger: {e2}") except Exception as e: - print(f"Error clearing Metric context: {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 + """ + Cancel the current task. + + Args: + internal: If True, suppress output messages + """ try: current_experiment = self.db.get_running_experiment() - except DBException: + except Exception: if not internal: print("No experiment is currently running. Nothing to cancel.") return - # create logger - logger = RFLogger().create_logger(current_experiment["experiment_name"]) + experiment_name = current_experiment["experiment_name"] + experiments_path = current_experiment.get("experiments_path", RF_EXPERIMENT_PATH) + + # Create logger + logger = RFLogger( + experiment_name=experiment_name, + experiment_path=experiments_path, + ).get_logger("ExperimentUtils") try: current_task = self.db.get_experiment_current_task() - except DBException: - msg = "No task is currently running. Nothing to cancel." + except Exception: if not internal: - print(msg) - logger.info(msg) + print("No task is currently running. Nothing to cancel.") return - # reset experiment states and set current task to idle + # 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") + logger.debug("Reset experiment states and set current experiment task to idle") def get_runs_info(self) -> pd.DataFrame: - """Get the run info""" + """ + Get run info for all runs in the experiment (fit mode). + + Returns: + DataFrame with run information + """ 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")} + 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 = run_details["config_leaf"].copy() if run_details["config_leaf"] else {} config_leaf.pop("additional_kwargs", None) new_run_details["config"] = config_leaf @@ -303,94 +322,87 @@ def get_runs_info(self) -> pd.DataFrame: 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. + except Exception as e: + raise ExperimentException(f"Error getting runs info: {e}") from e - Returns: - pd.DataFrame: A DataFrame containing all runs information with run_id as first column. + def _create_experiment_internal( + self, + given_name: str, + experiments_path: str, + ) -> tuple[int, str, str | None]: """ - 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] + Create new experiment - generate unique name and write to database. - # 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) + Args: + given_name: Desired experiment name + experiments_path: Path to experiments directory - # 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 | None]: - """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 - Returns: experiment_id, experiment_name, metric_experiment_id (or None) + Returns: + Tuple of (experiment_id, experiment_name, metric_experiment_id or None) """ 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()) + experiment_name = self._generate_unique_experiment_name( + given_name, + self.db.get_all_experiment_names(), + ) - # Create Metricexperiment only if available - metric_experiment_id = None - if RF_MLFLOW_ENABLED=="true": - import mlflow # Lazy import to avoid connection attempts in tensorboard-only mode + # Clear all non-experiment tables before creating new experiment + self.db.reset_all_tables(experiments_table=False) + # Create MLflow experiment if enabled + metric_experiment_id = None + if RF_MLFLOW_ENABLED == "true": try: - # create logger - logger = RFLogger().create_logger(experiment_name) - metric_logger_config = RFMetricLogger.get_default_metric_loggers(experiment_name=experiment_name) + import mlflow + + from rapidfireai.metrics import RFMetricLogger + + logger = RFLogger( + experiment_name=experiment_name, + experiment_path=experiments_path, + ).get_logger("ExperimentUtils") + + metric_logger_config = RFMetricLogger.get_default_metric_loggers( + experiment_name=experiment_name + ) metric_logger = RFMetricLogger(metric_logger_config, logger=logger) metric_experiment_id = metric_logger.create_experiment(experiment_name) mlflow.tracing.disable_notebook_display() except Exception as e: - # Catch MLflow-specific exceptions (mlflow.exceptions.RestException, etc.) raise ExperimentException(f"Error creating Metric experiment: {e}") from e - # write new experiment details to database + # Write new experiment to database experiment_id = self.db.create_experiment( - experiment_name, - metric_experiment_id, # Will be None for tensorboard-only - config_options={"experiments_path": experiments_path}, + experiment_name=experiment_name, + experiments_path=os.path.abspath(experiments_path), + metric_experiment_id=metric_experiment_id, + status=ExperimentStatus.RUNNING, ) + return experiment_id, experiment_name, metric_experiment_id + except ExperimentException: - # Re-raise ExperimentExceptions (including MLflow errors from above) raise except Exception as e: - # Catch any other unexpected errors raise ExperimentException(f"Error in _create_experiment_internal: {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""" + """ + Generate a unique experiment name by incrementing suffix if needed. + + Args: + name: Desired base name + existing_names: List of existing experiment names + + Returns: + Unique experiment name + """ if not name: name = "rf-exp" pattern = r"^(.+?)(_(\d+))?$" - max_attempts = 1000 # Prevent infinite loops + max_attempts = 1000 attempts = 0 new_name = name @@ -404,7 +416,6 @@ def _generate_unique_experiment_name(self, name: str, existing_names: list[str]) 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 @@ -418,3 +429,6 @@ def _generate_unique_experiment_name(self, name: str, existing_names: list[str]) raise ExperimentException("Could not generate unique experiment name") return new_name + + +__all__ = ["ExperimentUtils"] diff --git a/rapidfireai/utils/logging.py b/rapidfireai/utils/logging.py new file mode 100644 index 00000000..9726a744 --- /dev/null +++ b/rapidfireai/utils/logging.py @@ -0,0 +1,185 @@ +""" +RapidFire AI Logging Module + +Provides Python standard logging with Ray-compatible file handling. +""" + +import logging +import os +import threading +from pathlib import Path + +from rapidfireai.utils.constants import ( + RF_EXPERIMENT_PATH, + RF_LOG_FILENAME, + RF_LOG_PATH, + RF_TRAINING_LOG_FILENAME, + LogType, +) +from rapidfireai.utils.os_utils import mkdir_p + + +class AsyncioCleanupFilter(logging.Filter): + """Filter out harmless asyncio cleanup errors during shutdown.""" + + def filter(self, record): + msg = str(record.getMessage()) + # Suppress TCPTransport cleanup errors + if "TCPTransport" in msg and "closed=True" in msg: + return False + # Suppress Task exception errors for httpx/OpenAI client cleanup + if "Task exception was never retrieved" in msg: + if "AsyncClient.aclose" in msg: + return False + return True + + +# Third-party loggers to suppress +THIRD_PARTY_LOGGERS = [ + "ray", + "vllm", + "torch", + "transformers", + "datasets", + "huggingface_hub", + "bitsandbytes", + "peft", + "trl", + "accelerate", + "langchain", + "langchain_core", + "langchain_community", + "openai", + "httpx", + "urllib3", + "filelock", +] + + +class RFLogger: + """ + RapidFire logger using Python standard logging. + + Features: + - Ray worker detection (avoids file handler conflicts) + - Thread-safe initialization + - Third-party log suppression + """ + + # Class-level state shared across instances + _file_handlers: dict[str, logging.FileHandler] = {} + _initialized_experiments: set[str] = set() + _lock = threading.Lock() + + def __init__( + self, + experiment_name: str = "unknown", + experiment_path: str = RF_EXPERIMENT_PATH, + level: str = "INFO", + log_type: LogType = LogType.RF_LOG, + ): + self._experiment_name = experiment_name + self._experiment_path = experiment_path + self.level = level.upper() + self.log_type = log_type + + # Suppress third-party logs via environment variables + os.environ.setdefault("VLLM_LOGGING_LEVEL", "ERROR") + os.environ.setdefault("RAY_LOG_TO_STDERR", "0") + + # Only set up file handlers on main process (not Ray workers) and when experiment name is known + is_ray_worker = os.environ.get("RAY_WORKER_MODE") == "WORKER" + if not is_ray_worker and experiment_name != "unknown": + self._setup_file_handler() + + def _get_log_file_path(self) -> str: + """Get the log file path based on log type.""" + if self.log_type == LogType.TRAINING_LOG: + return os.path.join(RF_LOG_PATH, self._experiment_name, RF_TRAINING_LOG_FILENAME) + return os.path.join(RF_LOG_PATH, self._experiment_name, RF_LOG_FILENAME) + + def _setup_file_handler(self): + """Set up file handler for this logger (thread-safe).""" + with RFLogger._lock: + handler_key = f"{self.log_type.value}_{self._experiment_name}" + root_logger = logging.getLogger() + + # Check if handler already exists + if handler_key in RFLogger._file_handlers: + file_handler = RFLogger._file_handlers[handler_key] + # Ensure handler is attached to root logger (may have been removed during notebook re-runs) + if file_handler not in root_logger.handlers: + root_logger.addHandler(file_handler) + return + + # Create log directory + log_dir = Path(RF_LOG_PATH) / self._experiment_name + mkdir_p(log_dir.absolute()) + + # Create file handler + log_file_path = self._get_log_file_path() + log_format = "%(asctime)s | %(name)s | %(levelname)s | %(filename)s:%(lineno)d | %(message)s" + + file_handler = logging.FileHandler(log_file_path) + file_handler.setFormatter(logging.Formatter(log_format, datefmt="%Y-%m-%d %H:%M:%S")) + file_handler.setLevel(self.level) + + RFLogger._file_handlers[handler_key] = file_handler + + # Configure root logger once per experiment + if self._experiment_name not in RFLogger._initialized_experiments: + RFLogger._initialized_experiments.add(self._experiment_name) + + # Remove existing handlers (prevent console output) + for handler in root_logger.handlers[:]: + root_logger.removeHandler(handler) + + root_logger.setLevel(self.level) + root_logger.addFilter(AsyncioCleanupFilter()) + + # Suppress asyncio logger + asyncio_logger = logging.getLogger("asyncio") + asyncio_logger.addFilter(AsyncioCleanupFilter()) + asyncio_logger.setLevel(logging.CRITICAL) + + # Suppress third-party loggers + for logger_name in THIRD_PARTY_LOGGERS: + logging.getLogger(logger_name).setLevel(logging.CRITICAL) + logging.getLogger(logger_name).propagate = False + + # Add handler to root logger + root_logger.addHandler(file_handler) + + def get_logger(self, name: str = "unknown") -> logging.Logger: + """ + Get a configured logger instance. + + Args: + name: Name for the logger (e.g., 'controller', 'worker_0') + + Returns: + A standard Python logger + """ + logger = logging.getLogger(f"{self._experiment_name}:{name}") + logger.setLevel(self.level) + return logger + + +class TrainingLogger(RFLogger): + """Training-specific logger that writes to training.log instead of rapidfire.log.""" + + def __init__( + self, + experiment_name: str = "unknown", + experiment_path: str = RF_EXPERIMENT_PATH, + level: str = "DEBUG", + ): + super().__init__( + experiment_name=experiment_name, + experiment_path=experiment_path, + level=level, + log_type=LogType.TRAINING_LOG, + ) + + +__all__ = ["RFLogger", "TrainingLogger"] diff --git a/rapidfireai/utils/os_utils.py b/rapidfireai/utils/os_utils.py index 53524797..85b6cb12 100644 --- a/rapidfireai/utils/os_utils.py +++ b/rapidfireai/utils/os_utils.py @@ -1,17 +1,18 @@ """Utility functions for OS information.""" -import subprocess -import re import os +import re +import subprocess from pathlib import Path + def mkdir_p(path: str, parents: bool = True, exist_ok: bool = True, notify: bool = True): """Create a directory if it does not exist.""" if not os.path.exists(path): - Path(path).mkdir(parents=parents, exist_ok=exist_ok) - if notify: - print(f"Created directory: {path}") - return + Path(path).mkdir(parents=parents, exist_ok=exist_ok) + if notify: + print(f"Created directory: {path}") + return if not os.path.isdir(path): raise OSError(f"Path exist and is not a directory: {path}") return @@ -20,72 +21,54 @@ def mkdir_p(path: str, parents: bool = True, exist_ok: bool = True, notify: bool def get_os_package_installed(package_pattern: str): """Get list of installed packages matching a pattern.""" import distro + dist_id = distro.id() - + try: - if dist_id in ['ubuntu', 'debian']: + if dist_id in ["ubuntu", "debian"]: # Use dpkg-query for Debian-based result = subprocess.run( - ['dpkg-query', '-W', '-f=${Package}\n', package_pattern], - capture_output=True, - text=True, - check=False + ["dpkg-query", "-W", "-f=${Package}\n", package_pattern], capture_output=True, text=True, check=False ) if result.returncode == 0: - return [pkg.strip() for pkg in result.stdout.strip().split('\n') if pkg.strip()] + return [pkg.strip() for pkg in result.stdout.strip().split("\n") if pkg.strip()] return [] - - elif dist_id in ['rhel', 'centos', 'fedora', 'rocky', 'almalinux']: + + elif dist_id in ["rhel", "centos", "fedora", "rocky", "almalinux"]: # Use rpm for Red Hat-based - result = subprocess.run( - ['rpm', '-qa', package_pattern], - capture_output=True, - text=True, - check=False - ) + result = subprocess.run(["rpm", "-qa", package_pattern], capture_output=True, text=True, check=False) if result.returncode == 0: - return [pkg.strip() for pkg in result.stdout.strip().split('\n') if pkg.strip()] + return [pkg.strip() for pkg in result.stdout.strip().split("\n") if pkg.strip()] return [] - - elif dist_id in ['arch', 'manjaro']: + + elif dist_id in ["arch", "manjaro"]: # Use pacman for Arch-based - result = subprocess.run( - ['pacman', '-Qq'], - capture_output=True, - text=True, - check=False - ) + result = subprocess.run(["pacman", "-Qq"], capture_output=True, text=True, check=False) if result.returncode == 0: - all_packages = result.stdout.strip().split('\n') + all_packages = result.stdout.strip().split("\n") # Convert shell glob pattern to regex - pattern_regex = package_pattern.replace('*', '.*') + pattern_regex = package_pattern.replace("*", ".*") return [pkg for pkg in all_packages if re.match(pattern_regex, pkg)] return [] - - elif dist_id in ['opensuse', 'sles']: + + elif dist_id in ["opensuse", "sles"]: # Use rpm for openSUSE - result = subprocess.run( - ['rpm', '-qa', package_pattern], - capture_output=True, - text=True, - check=False - ) + result = subprocess.run(["rpm", "-qa", package_pattern], capture_output=True, text=True, check=False) if result.returncode == 0: - return [pkg.strip() for pkg in result.stdout.strip().split('\n') if pkg.strip()] + return [pkg.strip() for pkg in result.stdout.strip().split("\n") if pkg.strip()] return [] - + else: # Fallback: try dpkg first, then rpm - for cmd in [['dpkg-query', '-W', '-f=${Package}\n', package_pattern], - ['rpm', '-qa', package_pattern]]: + for cmd in [["dpkg-query", "-W", "-f=${Package}\n", package_pattern], ["rpm", "-qa", package_pattern]]: try: result = subprocess.run(cmd, capture_output=True, text=True, check=False) if result.returncode == 0 and result.stdout.strip(): - return [pkg.strip() for pkg in result.stdout.strip().split('\n') if pkg.strip()] + return [pkg.strip() for pkg in result.stdout.strip().split("\n") if pkg.strip()] except FileNotFoundError: continue return [] - + except Exception as e: print(f"Error checking packages: {e}") return [] diff --git a/rapidfireai/utils/serialize.py b/rapidfireai/utils/serialize.py new file mode 100644 index 00000000..1f4ff8c5 --- /dev/null +++ b/rapidfireai/utils/serialize.py @@ -0,0 +1,177 @@ +""" +RapidFire AI Serialization + +Serialization utilities for storing complex objects in the database. +""" + +import base64 +import json +from typing import Any + + +def encode_payload(payload: object) -> str: + """ + Encode a Python object for database storage using dill. + + Args: + payload: Any Python object (including functions, classes, etc.) + + Returns: + Base64-encoded string representation + """ + import dill + return base64.b64encode(dill.dumps(payload)).decode("utf-8") + + +def decode_db_payload(payload: str) -> object: + """ + Decode a database payload back to a Python object. + + Args: + payload: Base64-encoded string from database + + Returns: + Original Python object + """ + import dill + return dill.loads(base64.b64decode(payload)) + + +def extract_pipeline_config_json(pipeline_config: dict[str, Any]) -> dict[str, Any]: + """ + Extract JSON-serializable data from a pipeline config dictionary. + + Extracts only serializable parameters (dicts, strings, ints, etc.) and ignores + functions, classes, and other non-serializable objects. This is used for storing + a JSON representation in the database for analytics/display purposes. + + The actual pipeline config (with functions and classes) should be stored using + encode_payload/decode_db_payload in the pipeline_config column. + + Args: + pipeline_config: Pipeline config dict with keys: + - "pipeline": RFvLLMModelConfig or RFOpenAIAPIModelConfig instance + - "batch_size": int + - "preprocess_fn": function (skipped) + - "postprocess_fn": function (skipped) + - "compute_metrics_fn": function (skipped) + - "accumulate_metrics_fn": function (skipped) + - "online_strategy_kwargs": dict (optional) + + Returns: + Dictionary with only JSON-serializable data from the pipeline config + """ + # Handle None or empty config + if not pipeline_config: + return {} + + json_config = {} + + # Extract batch_size if present + if "batch_size" in pipeline_config: + json_config["batch_size"] = pipeline_config["batch_size"] + + # Extract online_strategy_kwargs if present + if "online_strategy_kwargs" in pipeline_config: + json_config["online_strategy_kwargs"] = pipeline_config["online_strategy_kwargs"] + + # Extract pipeline type and model-specific params + if "pipeline" in pipeline_config: + pipeline = pipeline_config["pipeline"] + + # Helper function to extract RAG params + def extract_rag_params(rag_spec): + """Extract RAG parameters from rag_spec similar to controller logic.""" + if rag_spec is None: + return None + + rag_config = {} + rag_config["search_type"] = getattr(rag_spec, "search_type", None) + + if hasattr(rag_spec, "search_kwargs") and rag_spec.search_kwargs is not None: + rag_config["k"] = rag_spec.search_kwargs.get("k", None) + + if hasattr(rag_spec, "reranker_kwargs") and rag_spec.reranker_kwargs is not None: + rag_config["top_n"] = rag_spec.reranker_kwargs.get("top_n", None) + + if hasattr(rag_spec, "text_splitter") and rag_spec.text_splitter is not None: + rag_config["chunk_size"] = getattr(rag_spec.text_splitter, "_chunk_size", None) + rag_config["chunk_overlap"] = getattr(rag_spec.text_splitter, "_chunk_overlap", None) + + # Only return rag_config if it has at least one non-None value + filtered_rag_config = {k: v for k, v in rag_config.items() if v is not None} + return filtered_rag_config if filtered_rag_config else None + + # Try to import the model config classes for type checking + try: + from rapidfireai.automl.model_config import RFOpenAIAPIModelConfig, RFvLLMModelConfig + + if isinstance(pipeline, RFvLLMModelConfig): + json_config["pipeline_type"] = "vllm" + + # Extract model_config (dict) + if hasattr(pipeline, "model_config") and pipeline.model_config is not None: + json_config["model_config"] = pipeline.model_config + + # Extract sampling_params from _user_params (original dict, not SamplingParams object) + if hasattr(pipeline, "_user_params") and "sampling_params" in pipeline._user_params: + json_config["sampling_params"] = pipeline._user_params["sampling_params"] + + # Extract RAG params if present + if hasattr(pipeline, "rag") and pipeline.rag is not None: + rag_config = extract_rag_params(pipeline.rag) + if rag_config: + json_config["rag_config"] = rag_config + + elif isinstance(pipeline, RFOpenAIAPIModelConfig): + json_config["pipeline_type"] = "openai" + + # Extract client_config (dict) - filter out sensitive keys + if hasattr(pipeline, "client_config") and pipeline.client_config is not None: + sensitive_keys = {"api_key", "secret", "token", "password", "key"} + json_config["client_config"] = { + k: v for k, v in pipeline.client_config.items() + if k.lower() not in sensitive_keys + } + + # Extract model_config (dict) + if hasattr(pipeline, "model_config") and pipeline.model_config is not None: + json_config["model_config"] = pipeline.model_config + + # Extract sampling_params using sampling_params_to_dict (extracts from model_config) + if hasattr(pipeline, "sampling_params") and pipeline.sampling_params is not None: + json_config["sampling_params"] = pipeline.sampling_params_to_dict() + + # Extract rate limiting params + if hasattr(pipeline, "rpm_limit") and pipeline.rpm_limit is not None: + json_config["rpm_limit"] = pipeline.rpm_limit + if hasattr(pipeline, "tpm_limit") and pipeline.tpm_limit is not None: + json_config["tpm_limit"] = pipeline.tpm_limit + if hasattr(pipeline, "max_completion_tokens") and pipeline.max_completion_tokens is not None: + json_config["max_completion_tokens"] = pipeline.max_completion_tokens + + # Extract RAG params if present + if hasattr(pipeline, "rag") and pipeline.rag is not None: + rag_config = extract_rag_params(pipeline.rag) + if rag_config: + json_config["rag_config"] = rag_config + + except ImportError: + # Model config classes not available, try generic extraction + if hasattr(pipeline, "__class__"): + class_name = pipeline.__class__.__name__.lower() + if "vllm" in class_name: + json_config["pipeline_type"] = "vllm" + elif "openai" in class_name: + json_config["pipeline_type"] = "openai" + + # Validate JSON serializability + try: + json.dumps(json_config) + except (TypeError, ValueError) as e: + raise ValueError(f"Failed to serialize pipeline config to JSON: {e}") from e + + return json_config + + +__all__ = ["encode_payload", "decode_db_payload", "extract_pipeline_config_json"] diff --git a/rapidfireai/version.py b/rapidfireai/version.py index 2c5879da..e636a424 100644 --- a/rapidfireai/version.py +++ b/rapidfireai/version.py @@ -3,4 +3,4 @@ """ __version__ = "0.15.2" -__version_info__ = (0, 15, 2) \ No newline at end of file +__version_info__ = (0, 15, 2) diff --git a/requirements.txt b/requirements.txt index 29ad5178..3d50e319 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,7 @@ # ml pandas>=2.0.0,<2.3.0 # Colab compatibility (2.2.2) and cudf constraint torch>=2.7.0 -transformers>=4.55.2 +transformers>=4.55.2,<5.0.0 peft>=0.17.0 trl>=0.21.0 bitsandbytes>=0.47.0 @@ -12,6 +12,9 @@ evaluate>=0.4.5 rouge-score>=0.1.2 sentencepiece>=0.2.1 +# distributed computing +ray[default]>=2.46.0 # Ray for distributed training (fit and evals) + # misc dill>=0.3.8 mlflow>=3.2.0 @@ -19,7 +22,6 @@ pytest>=8.4.1 flask>=3.1.1 flask-cors>=6.0.1 requests>=2.32.0 # Relaxed for Colab (2.32.4) -loguru>=0.7.3 ipython>=7.34.0 # Colab compatibility (7.34.0) jupyter>=1.1.1 ipywidgets>=7.3.4,<9.0.0 diff --git a/setup/evals/1_setup_conda.sh b/setup/evals/1_setup_conda.sh deleted file mode 100755 index eb36bb2b..00000000 --- a/setup/evals/1_setup_conda.sh +++ /dev/null @@ -1,51 +0,0 @@ -#!/bin/bash - -# Miniconda Installation and Python 3.10 Setup Script -# This script installs miniconda and sets up Python 3.10 in the base environment - -echo "Starting Miniconda installation and Python 3.10 setup..." - -# Step 1: Check if conda is already installed -if command -v conda &> /dev/null; then - echo "WARNING: Conda is already installed. Skipping installation..." - echo "Current conda version: $(conda --version)" -else - echo "Downloading and installing Miniconda..." - # Download Miniconda installer - wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh - - # Install Miniconda in batch mode - bash ~/miniconda.sh -b -p $HOME/miniconda3 - - # Initialize conda for bash - $HOME/miniconda3/bin/conda init bash - - # Remove installer file - rm ~/miniconda.sh - - # Source bashrc to activate conda automatically - source ~/.bashrc - - # Also source conda profile directly to ensure it's available - source $HOME/miniconda3/etc/profile.d/conda.sh -fi - -# Check if conda is available, if not source it -if ! command -v conda &> /dev/null; then - echo "Conda not found in PATH, sourcing conda setup..." - source $HOME/miniconda3/etc/profile.d/conda.sh -fi - -# Accept Terms of Service for conda channels -echo "Accepting conda Terms of Service..." -conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main -conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r - -# Install Python 3.10 in base environment -echo "Installing Python 3.10 in base environment..." -conda install python=3.10 -y - -# Verify Python version -python --version - -echo "Setup complete! Python 3.10 is now installed in your conda base environment." \ No newline at end of file diff --git a/setup/evals/2_setup_conda_env.sh b/setup/evals/2_setup_conda_env.sh deleted file mode 100755 index e15f7091..00000000 --- a/setup/evals/2_setup_conda_env.sh +++ /dev/null @@ -1,58 +0,0 @@ -#!/bin/bash - -# Configuration -ENV_NAME="infer" - -# Exit on any error -set -e - -echo "Setting up conda environment '$ENV_NAME' with Python 3.10..." - -# Source conda if not already available -if ! command -v conda &> /dev/null; then - if [ -f "$HOME/miniconda3/etc/profile.d/conda.sh" ]; then - source "$HOME/miniconda3/etc/profile.d/conda.sh" - else - echo "ERROR: Conda not found. Please run 1_setup_conda.sh first." - exit 1 - fi -fi - -# Accept conda terms of service -conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main -conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r - -# Create conda environment with Python 3.10 -conda create -n $ENV_NAME python=3.10 -y - -# Activate the environment -source $(conda info --base)/etc/profile.d/conda.sh -conda activate $ENV_NAME - -echo "Installing required packages..." -pip install torch==2.5.1 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124 - -# Install uv package manager for faster installations -pip install uv - -# Install flashinfer-python -pip install flashinfer-python -i https://flashinfer.ai/whl/cu124/torch2.5/ - -# Install ML/AI packages using uv for speed -uv pip install ray -uv pip install gpustat -uv pip install vllm==0.7.2 --torch-backend=cu124 -uv pip install tensorboard -uv pip install ipywidgets ipykernel pandas pyarrow numpy datasets==3.6.0 -uv pip install langchain langchain-core langchain-community unstructured langchain-openai langchain-huggingface faiss-gpu - -# Install flash-attention -pip install flash-attn --no-build-isolation - -# Install rf-inferno package in editable mode -echo "Installing rf-inferno package in editable mode..." -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -PROJECT_ROOT="$(dirname "$SCRIPT_DIR")" -pip install -e "$PROJECT_ROOT" - -echo "Environment setup complete!" \ No newline at end of file diff --git a/setup/evals/requirements-local.txt b/setup/evals/requirements-local.txt deleted file mode 100644 index b64fa275..00000000 --- a/setup/evals/requirements-local.txt +++ /dev/null @@ -1,45 +0,0 @@ -# Non Colab Packages -# evals dependencies -# Core ML/AI Framework -torch<=2.8.0 -sentence-transformers>=5.1.0 - -# Distributed Computing -ray>=2.49.2 -ray[default]>=2.49.2 - -# LLM Inference -transformers>=4.56.1,<5.0.0 - -# OpenAI API -openai>=1.106.1 -tiktoken>=0.12.0 - -# LangChain Ecosystem -langchain>=1.0.5 -langchain-classic>=1.0.0 -langchain-core>=1.0.4 -langchain-community>=0.4.1 -langchain-openai>=1.0.2 -langchain-huggingface>=1.0.0 -langchain-pinecone>=0.2.13 -langchain-postgres>=0.0.17 - -# Data Manipulation & Display -unstructured>=0.18.15 - -# Other -requests==2.32.5 -datasets -jupyter -grpcio -faiss-gpu-cu12 - -# MLflow Dashboard dependencies -mlflow>=3.2.0 -gunicorn>=23.0.0 -flask-cors>=5.0.1 -loguru - -numpy==2.0.1 -protobuf<6.0.0 \ No newline at end of file diff --git a/setup/evals/requirements.txt b/setup/evals/requirements.txt deleted file mode 100644 index b13a4f5d..00000000 --- a/setup/evals/requirements.txt +++ /dev/null @@ -1,50 +0,0 @@ -# Core ML/AI Framework -torch==2.5.1 -torchvision -torchaudio -transformers -datasets==3.6.0 -sentence-transformers - -# Distributed Computing -ray - -# LLM Inference -vllm==0.7.2 -flash-attn -flashinfer-python - -# OpenAI API -openai -tiktoken - -# LangChain Ecosystem -langchain -langchain-core -langchain-community -langchain-openai -langchain-huggingface - -# Vector Search -faiss-gpu - -# Statistical Analysis -scipy - -# Data Manipulation & Display -pandas - -# REST API (Dispatcher) -flask -flask-cors -waitress - -# Notebook -ipykernel -ipywidgets -tensorboard - -jq - -# Logging -loguru \ No newline at end of file diff --git a/setup/evals/setup.sh b/setup/evals/setup.sh deleted file mode 100755 index 8519b996..00000000 --- a/setup/evals/setup.sh +++ /dev/null @@ -1,54 +0,0 @@ -#!/bin/bash - -# Complete Conda Environment Setup Script -# This script runs both conda installation and environment setup in sequence - -set -e # Exit on any error - -echo "Starting complete conda and environment setup..." -echo "" - -# Get the directory where this script is located -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" - -# Set the experiment path -export RF_EXPERIMENT_PATH="${SCRIPT_DIR}/../rapidfire_experiments" - -# Step 1: Run conda installation script -echo "Step 1: Installing Miniconda and Python 3.10..." -echo "----------------------------------------" -bash "${SCRIPT_DIR}/1_setup_conda.sh" - -echo "" -echo "Conda installation completed!" -echo "" - -# Source conda to make it available in this shell -if [ -f "$HOME/miniconda3/etc/profile.d/conda.sh" ]; then - source "$HOME/miniconda3/etc/profile.d/conda.sh" -fi - -# Step 2: Run environment setup script -echo "Step 2: Setting up 'infer' conda environment..." -echo "----------------------------------------" -bash "${SCRIPT_DIR}/2_setup_conda_env.sh" - -# Step 3: Install requirements from requirements.txt -echo "" -echo "Step 3: Installing requirements from requirements.txt..." -echo "----------------------------------------" -conda run -n infer pip3 install -r "${SCRIPT_DIR}/../requirements.txt" - -echo "" -echo "Complete setup finished!" -echo "" -echo "Summary:" -echo " - Miniconda installed with Python 3.10" -echo " - 'infer' conda environment created" -echo " - PyTorch 2.5.1 with CUDA 12.4 installed" -echo " - All ML/AI packages installed" -echo " - Requirements from requirements.txt installed" -echo "" -echo "To get started:" -echo " conda activate infer" -echo "" diff --git a/setup/fit/1_setup_miniconda_python312.sh b/setup/fit/1_setup_miniconda_python312.sh deleted file mode 100644 index bcb6b2ea..00000000 --- a/setup/fit/1_setup_miniconda_python312.sh +++ /dev/null @@ -1,61 +0,0 @@ -#!/bin/bash - -# Miniconda Installation and Python 3.12 Setup Script -# This script installs miniconda and sets up Python 3.12 in the base environment - -echo "๐Ÿ Starting Miniconda installation and Python 3.12 setup..." - -# Step 1: Check if conda is already installed -if command -v conda &> /dev/null; then - echo "โš ๏ธ Conda is already installed. Skipping installation..." - echo "Current conda version: $(conda --version)" -else - echo "๐Ÿ“ฅ Downloading Miniconda installer..." - # Step 2: Download Miniconda installer - wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh - - echo "๐Ÿ”ง Installing Miniconda..." - # Step 3: Install Miniconda in batch mode - bash ~/miniconda.sh -b -p $HOME/miniconda3 - - echo "โš™๏ธ Initializing conda for bash..." - # Step 4: Initialize conda for bash - $HOME/miniconda3/bin/conda init bash - - echo "๐Ÿงน Cleaning up installer file..." - # Step 5: Remove installer file - rm ~/miniconda.sh - - echo "๐Ÿ”„ Sourcing bashrc to activate conda..." - # Step 6: Source bashrc to activate conda (this needs to be done manually) - echo "โš ๏ธ Please run 'source ~/.bashrc' or restart your terminal to activate conda" -fi - -# Step 7: Accept Terms of Service for conda channels -echo "๐Ÿ“‹ Accepting conda Terms of Service..." -conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main -conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r - -# Step 8: Install Python 3.12 in base environment -echo "๐Ÿ Installing Python 3.12 in base environment..." -conda install python=3.12 -y - -# Step 9: Verify Python version -echo "โœ… Installation complete! Verifying Python version..." -python --version - -echo "๐ŸŽ‰ Setup complete! Python 3.12 is now installed in your conda base environment." -echo "" -echo "๐Ÿ”ง Manual steps required:" -echo " 1. If conda was just installed, run: source ~/.bashrc" -echo " 2. Activate base environment: conda activate base" -echo "" -echo "๐Ÿ“ Commands that were executed:" -echo " - wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh" -echo " - bash ~/miniconda.sh -b -p \$HOME/miniconda3" -echo " - \$HOME/miniconda3/bin/conda init bash" -echo " - source ~/.bashrc" -echo " - conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main" -echo " - conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r" -echo " - conda install python=3.12 -y" -echo " - rm ~/miniconda.sh" \ No newline at end of file diff --git a/setup/fit/2_install_ml_packages.sh b/setup/fit/2_install_ml_packages.sh deleted file mode 100644 index daaeefdd..00000000 --- a/setup/fit/2_install_ml_packages.sh +++ /dev/null @@ -1,117 +0,0 @@ -#!/bin/bash - -# ML Packages Installation and Jupyter Kernel Setup Script -# This script installs ML/AI packages and creates a Jupyter kernel for the conda base environment - -echo "๐Ÿ“ฆ Starting ML packages installation and Jupyter kernel setup..." - -# Check if we're in a conda environment -if [[ -z "${CONDA_DEFAULT_ENV}" ]]; then - echo "โš ๏ธ Warning: No conda environment detected. Please activate conda first:" - echo " Run: conda activate base" - exit 1 -fi - -echo "โœ… Conda environment detected: ${CONDA_DEFAULT_ENV}" - -# Check if we're in the correct directory (should have requirements.txt) -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -REQUIREMENTS_FILE="${SCRIPT_DIR}/requirements.txt" - -if [[ ! -f "$REQUIREMENTS_FILE" ]]; then - echo "โŒ Error: requirements.txt not found in ${SCRIPT_DIR}" - echo " Please ensure requirements.txt exists in the same directory as this script." - exit 1 -fi - -echo "๐Ÿ“‹ Found requirements.txt at: $REQUIREMENTS_FILE" - -# Update pip to latest version -echo "๐Ÿ”„ Updating pip to latest version..." -pip install --upgrade pip - -# Install packages from requirements.txt -echo "๐Ÿ“ฆ Installing packages from requirements.txt..." -echo " This may take several minutes..." -pip install -r "$REQUIREMENTS_FILE" - -# Check if installation was successful -if [[ $? -ne 0 ]]; then - echo "โŒ Error: Package installation failed!" - echo " Please check the error messages above and try again." - exit 1 -fi - -echo "โœ… Package installation completed successfully!" - -# Install/update ipykernel for the current environment -echo "๐Ÿ”ง Setting up Jupyter kernel for conda base environment..." - -# Get the current conda environment name and python path -CURRENT_ENV_NAME="${CONDA_DEFAULT_ENV}" -PYTHON_PATH=$(which python) - -echo " Environment: $CURRENT_ENV_NAME" -echo " Python path: $PYTHON_PATH" - -# Install the kernel -python -m ipykernel install --user --name="$CURRENT_ENV_NAME" --display-name="Python ($CURRENT_ENV_NAME)" - -# Check if kernel installation was successful -if [[ $? -eq 0 ]]; then - echo "โœ… Jupyter kernel installed successfully!" -else - echo "โŒ Error: Jupyter kernel installation failed!" - exit 1 -fi - -# List available kernels -echo "๐Ÿ“‹ Available Jupyter kernels:" -jupyter kernelspec list - -# Verify key packages -echo "๐Ÿ” Verifying key package installations..." -python -c " -import sys -packages_to_check = [ - 'torch', 'transformers', 'peft', 'bitsandbytes', - 'datasets', 'huggingface_hub', 'jupyter', 'ipykernel' -] - -print('Package verification:') -for package in packages_to_check: - try: - __import__(package) - print(f' โœ… {package}') - except ImportError: - print(f' โŒ {package} - NOT FOUND') - sys.exit(1) - -print('\\n๐ŸŽ‰ All key packages verified successfully!') -" - -if [[ $? -ne 0 ]]; then - echo "โŒ Some packages failed verification. Please check the installation." - exit 1 -fi - -echo "" -echo "๐ŸŽ‰ Setup complete! Your ML development environment is ready." -echo "" -echo "๐Ÿš€ Next steps:" -echo " 1. Start Jupyter Lab: jupyter lab" -echo " 2. Or start Jupyter Notebook: jupyter notebook" -echo " 3. Create a new notebook and select the kernel: Python ($CURRENT_ENV_NAME)" -echo "" -echo "๐Ÿ“ What was installed:" -echo " โ€ข Core ML libraries: PyTorch, Transformers, PEFT, BitsAndBytes" -echo " โ€ข HuggingFace ecosystem: Datasets, Hub, Accelerate, Evaluate" -echo " โ€ข Jupyter ecosystem: JupyterLab, Notebook, IPykernel, IPywidgets" -echo " โ€ข Data science tools: NumPy, Pandas, Matplotlib, Seaborn, Plotly" -echo " โ€ข Utilities: tqdm, rich, scikit-learn, scipy" -echo " โ€ข Optional: Weights & Biases, TensorBoard" -echo "" -echo "๐Ÿ“š Useful commands:" -echo " โ€ข Check installed packages: pip list" -echo " โ€ข Update a package: pip install --upgrade " -echo " โ€ข List Jupyter kernels: jupyter kernelspec list" \ No newline at end of file diff --git a/setup/fit/README.md b/setup/fit/README.md deleted file mode 100644 index 6288d01e..00000000 --- a/setup/fit/README.md +++ /dev/null @@ -1,108 +0,0 @@ -# Setup Scripts for ML/AI Development Environment - -This directory contains scripts to set up a complete ML/AI development environment with Miniconda, Python 3.12, and essential packages for machine learning and data science. - -## ๐Ÿš€ Quick Start - -### Step 1: Install Miniconda and Python 3.12 -```bash -./setup_miniconda_python312.sh -``` - -### Step 2: Install ML Packages and Setup Jupyter -```bash -# Make sure conda base environment is active -conda activate base - -# Install packages and setup Jupyter kernel -./install_ml_packages.sh -``` - -### Step 3: Start Jupyter -```bash -# Option 1: JupyterLab (recommended) -jupyter lab - -# Option 2: Classic Jupyter Notebook -jupyter notebook -``` - -## ๐Ÿ“ Files Description - -### `setup_miniconda_python312.sh` -- Downloads and installs Miniconda -- Sets up Python 3.12 in the base environment -- Handles conda initialization and Terms of Service acceptance -- Includes error checking and informative output - -### `requirements.txt` -Contains all necessary packages for ML/AI development: -- **Core ML**: PyTorch, Transformers, PEFT, BitsAndBytes -- **HuggingFace**: Datasets, Hub, Accelerate, Evaluate -- **Jupyter**: JupyterLab, Notebook, IPykernel, IPywidgets -- **Data Science**: NumPy, Pandas, Matplotlib, Seaborn, Plotly -- **Utilities**: tqdm, rich, scikit-learn, scipy -- **Optional**: Weights & Biases, TensorBoard - -### `install_ml_packages.sh` -- Installs all packages from requirements.txt -- Creates a Jupyter kernel for the conda base environment -- Verifies package installations -- Provides usage instructions - -## ๐Ÿ”ง Manual Setup Alternative - -If you prefer to run commands manually: - -```bash -# 1. Install Miniconda -wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh -bash ~/miniconda.sh -b -p $HOME/miniconda3 -$HOME/miniconda3/bin/conda init bash -source ~/.bashrc - -# 2. Accept Terms of Service -conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main -conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r - -# 3. Install Python 3.12 -conda install python=3.12 -y - -# 4. Install packages -pip install -r requirements.txt - -# 5. Setup Jupyter kernel -python -m ipykernel install --user --name="base" --display-name="Python (base)" -``` - -## ๐ŸŽฏ What You Get - -After running these scripts, you'll have: -- โœ… Miniconda with Python 3.12 -- โœ… All essential ML/AI packages installed -- โœ… Jupyter environment ready to use -- โœ… Proper kernel configuration for notebooks -- โœ… Progress bars, visualization tools, and utilities - -## ๐Ÿš€ Next Steps - -1. Start JupyterLab: `jupyter lab` -2. Create a new notebook -3. Select the "Python (base)" kernel -4. Start coding! - -## ๐Ÿ“š Useful Commands - -```bash -# Check installed packages -pip list - -# Update a package -pip install --upgrade package_name - -# List available Jupyter kernels -jupyter kernelspec list - -# Check conda environment info -conda info -``` \ No newline at end of file diff --git a/setup/fit/build_and_install.sh b/setup/fit/build_and_install.sh deleted file mode 100644 index 0b552cad..00000000 --- a/setup/fit/build_and_install.sh +++ /dev/null @@ -1,88 +0,0 @@ -#!/bin/bash - -# RapidFire AI Build and Install Script -# This script builds and installs the rapidfireai package - -set -e # Exit on any error - -# Colors for output -RED='\033[0;31m' -GREEN='\033[0;32m' -YELLOW='\033[1;33m' -BLUE='\033[0;34m' -NC='\033[0m' # No Color - -# Function to print colored output -print_status() { - echo -e "${BLUE}[$(date '+%Y-%m-%d %H:%M:%S')]${NC} $1" -} - -print_success() { - echo -e "${GREEN}[$(date '+%Y-%m-%d %H:%M:%S')]${NC} $1" -} - -print_error() { - echo -e "${RED}[$(date '+%Y-%m-%d %H:%M:%S')]${NC} $1" -} - -print_warning() { - echo -e "${YELLOW}[$(date '+%Y-%m-%d %H:%M:%S')]${NC} $1" -} - -# Get the project root directory (parent of setup/fit/) -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" - -# Change to project root directory -cd "$PROJECT_ROOT" - -# Check if we're in the right directory -if [[ ! -f "pyproject.toml" ]]; then - print_error "This script must be run from the rapidfireai directory (pyproject.toml not found)" - exit 1 -fi - -# Check if we're in a virtual environment -if [[ -z "$VIRTUAL_ENV" ]] && [[ -z "$CONDA_DEFAULT_ENV" ]]; then - print_warning "Not in a virtual environment. This may cause permission issues." - print_status "Consider activating a virtual environment first:" - print_status " python3 -m venv .venv && source .venv/bin/activate" - read -p "Continue anyway? (y/N): " -n 1 -r - echo - if [[ ! $REPLY =~ ^[Yy]$ ]]; then - exit 1 - fi -fi - -# Clean previous builds -print_status "Cleaning previous builds..." -rm -rf build/ dist/ *.egg-info/ .eggs/ - -# Install build dependencies -print_status "Installing build dependencies..." -pip install --upgrade pip setuptools wheel build - -# Build the package -print_status "Building the package..." -python -m build - -# Install the package -print_status "Installing the package..." -pip install dist/*.whl - -print_success "Package installed successfully!" -print_status "You can now use the 'rapidfireai' command:" -print_status " rapidfireai --help" -print_status " rapidfireai start" -print_status " rapidfireai status" -print_status " rapidfireai stop" - -# Test the installation -print_status "Testing the installation..." -if command -v rapidfireai &> /dev/null; then - print_success "rapidfireai command is available" - rapidfireai --version -else - print_error "rapidfireai command not found in PATH" - print_status "You may need to restart your shell or activate your virtual environment" -fi \ No newline at end of file diff --git a/setup/fit/requirements-colab.txt b/setup/fit/requirements-colab.txt deleted file mode 100644 index c4a1e192..00000000 --- a/setup/fit/requirements-colab.txt +++ /dev/null @@ -1,22 +0,0 @@ -# Colab Packages -# fit dependencies -pandas>=2.0.0,<2.3.0 # Colab compatibility (2.2.2) and cudf constraint -psutil>=5.9.0 -torch>=2.8.0 -transformers>=4.55.2 -peft>=0.17.0 -trl==0.21.0 -bitsandbytes>=0.47.0 -nltk>=3.9.1 -evaluate>=0.4.5 -rouge-score>=0.1.2 -sentencepiece>=0.2.1 -mlflow>=3.2.0 -pytest>=8.4.1 -requests>=2.32.0 # Relaxed for Colab (2.32.4) -loguru>=0.7.3 -ipython>=7.34.0 # Colab compatibility (7.34.0) -jupyter>=1.1.1 -ipywidgets>=7.3.4,<9.0.0 # Support both v7 (Colab) and v8 (Jupyter) -tensorboard>=2.11.0 - diff --git a/setup/fit/requirements-local.txt b/setup/fit/requirements-local.txt deleted file mode 100644 index d7d6a196..00000000 --- a/setup/fit/requirements-local.txt +++ /dev/null @@ -1,64 +0,0 @@ -# Non Colab Packages -# fit dependencies -# Core ML/AI Framework -torch>=2.5.1 -torchvision>=0.20.1 -torchaudio>=2.5.1 -datasets>=3.6.0 -sentence-transformers==5.1.0 -gpustat==1.1.1 - -# Distributed Computing -click<8.3.0 -ray<=2.49.0 - -# LLM Inference -transformers>=4.56.1 - -# OpenAI API -openai==1.106.1 -tiktoken==0.11.0 - -# LangChain Ecosystem -langchain -langchain-classic -langchain-core -langchain-community -langchain-openai -langchain-huggingface - -# Vector Search -faiss-gpu-cu12==1.13.0 - -# Statistical Analysis -scipy>=1.16.1 - -# Data Manipulation & Display -pandas>=2.3.2 -pyarrow==21.0.0 -numpy>=1.26.4,<2.3 -unstructured==0.18.15 - - -# Notebook -ipykernel==6.30.1 -ipywidgets>=7.3.4,<9.0.0 # Support both v7 (Colab) and v8 (Jupyter) -tensorboard>=2.11.0 - - -# Other -psutil==7.0.0 -tqdm==4.67.1 -typing-extensions>=4.0.0 -peft>=0.17.0 -trl==0.21.0 -bitsandbytes>=0.47.0 -nltk>=3.9.1 -evaluate>=0.4.5 -rouge-score>=0.1.2 -sentencepiece>=0.2.1 -# mlflow>=3.2.0 -requests>=2.32.0 # Relaxed for Colab (2.32.4) -loguru>=0.7.3 -ipython>=7.34.0 # Colab compatibility (7.34.0) -jupyter>=1.1.1 diff --git a/setup/fit/requirements.txt b/setup/fit/requirements.txt deleted file mode 100644 index 625f47ab..00000000 --- a/setup/fit/requirements.txt +++ /dev/null @@ -1,169 +0,0 @@ -absl-py==2.3.1 -accelerate==1.9.0 -aiohappyeyeballs==2.6.1 -aiohttp==3.12.14 -aiosignal==1.4.0 -anyio==4.9.0 -argon2-cffi==25.1.0 -argon2-cffi-bindings==21.2.0 -arrow==1.3.0 -asttokens==3.0.0 -async-lru==2.0.5 -attrs==25.3.0 -babel==2.17.0 -beautifulsoup4==4.13.4 -bitsandbytes==0.46.1 -bleach==6.2.0 -blessed==1.21.0 -click==8.2.1 -comm==0.2.3 -contourpy==1.3.3 -cycler==0.12.1 -datasets==4.0.0 -debugpy==1.8.15 -decorator==5.2.1 -defusedxml==0.7.1 -dill==0.3.8 -einops==0.8.1 -evaluate==0.4.5 -executing==2.2.0 -fastjsonschema==2.21.1 -filelock==3.18.0 -fonttools==4.59.0 -fqdn==1.5.1 -frozenlist==1.7.0 -fsspec==2025.3.0 -gitdb==4.0.12 -GitPython==3.1.45 -gpustat==1.1.1 -grpcio==1.74.0 -h11==0.16.0 -hf-xet==1.1.5 -httpcore==1.0.9 -httpx==0.28.1 -huggingface-hub==0.34.2 -ipykernel==6.30.0 -ipython==9.4.0 -ipython_pygments_lexers==1.1.1 -ipywidgets==8.1.7 -isoduration==20.11.0 -jedi==0.19.2 -Jinja2==3.1.6 -joblib==1.5.1 -json5==0.12.0 -jsonpointer==2.1 -jsonschema==4.25.0 -jsonschema-specifications==2025.4.1 -jupyter==1.1.1 -jupyter-console==6.6.3 -jupyter-events==0.12.0 -jupyter-lsp==2.2.6 -jupyter_client==8.6.3 -jupyter_core==5.8.1 -jupyter_server==2.16.0 -jupyter_server_terminals==0.5.3 -jupyterlab==4.4.5 -jupyterlab_pygments==0.3.0 -jupyterlab_server==2.27.3 -jupyterlab_widgets==3.0.15 -kiwisolver==1.4.8 -lark==1.2.2 -Markdown==3.8.2 -MarkupSafe==3.0.2 -matplotlib==3.10.3 -matplotlib-inline==0.1.7 -mistune==3.1.3 -mpmath==1.3.0 -multidict==6.6.3 -multiprocess==0.70.16 -narwhals==2.0.0 -nbclient==0.10.2 -nbconvert==7.16.6 -nbformat==5.10.4 -nest-asyncio==1.6.0 -networkx==3.5 -notebook==7.4.4 -notebook_shim==0.2.4 -numpy==2.3.2 -nvidia-cublas-cu12==12.6.4.1 -nvidia-cuda-cupti-cu12==12.6.80 -nvidia-cuda-nvrtc-cu12==12.6.77 -nvidia-cuda-runtime-cu12==12.6.77 -nvidia-cudnn-cu12==9.5.1.17 -nvidia-cufft-cu12==11.3.0.4 -nvidia-cufile-cu12==1.11.1.6 -nvidia-curand-cu12==10.3.7.77 -nvidia-cusolver-cu12==11.7.1.2 -nvidia-cusparse-cu12==12.5.4.2 -nvidia-cusparselt-cu12==0.6.3 -nvidia-ml-py==12.575.51 -nvidia-nccl-cu12==2.26.2 -nvidia-nvjitlink-cu12==12.6.85 -nvidia-nvtx-cu12==12.6.77 -overrides==7.7.0 -pandas==2.3.1 -pandocfilters==1.5.1 -parso==0.8.4 -peft==0.16.0 -pexpect==4.9.0 -pillow==11.3.0 -plotly==6.2.0 -prometheus_client==0.22.1 -prompt_toolkit==3.0.51 -propcache==0.3.2 -protobuf==6.31.1 -psutil==7.0.0 -ptyprocess==0.7.0 -pure_eval==0.2.3 -pyarrow==21.0.0 -pyparsing==3.2.3 -python-dateutil==2.9.0.post0 -python-json-logger==3.3.0 -pytz==2025.2 -PyYAML==6.0.2 -pyzmq==27.0.0 -referencing==0.36.2 -regex==2024.11.6 -rfc3339-validator==0.1.4 -rfc3986-validator==0.1.1 -rfc3987-syntax==1.1.0 -rpds-py==0.26.0 -safetensors==0.5.3 -scikit-learn==1.7.1 -scipy==1.16.1 -seaborn==0.13.2 -Send2Trash==1.8.3 -sentry-sdk==2.33.2 -setuptools==78.1.1 -six==1.17.0 -smmap==5.0.2 -sniffio==1.3.1 -soupsieve==2.7 -stack-data==0.6.3 -sympy==1.14.0 -tensorboard==2.20.0 -tensorboard-data-server==0.7.2 -terminado==0.18.1 -threadpoolctl==3.6.0 -tinycss2==1.4.0 -tokenizers==0.21.4 -torch==2.7.1 -tornado==6.5.1 -trackio -traitlets==5.14.3 -transformers==4.54.0 -triton==3.3.1 -trl==0.19.1 -types-python-dateutil==2.9.0.20250708 -tzdata==2025.2 -uri-template==1.3.0 -wandb==0.21.0 -wcwidth==0.2.13 -webcolors==24.11.1 -webencodings==0.5.1 -websocket-client==1.8.0 -Werkzeug==3.1.3 -wheel==0.45.1 -widgetsnbextension==4.0.14 -xxhash==3.5.0 -yarl==1.20.1 diff --git a/setup/fit/start_dev.sh b/setup/fit/start_dev.sh deleted file mode 100755 index a33782d4..00000000 --- a/setup/fit/start_dev.sh +++ /dev/null @@ -1,504 +0,0 @@ -#!/bin/bash - -# RapidFire AI Multi-Service Startup Script -# This script starts MLflow server, API server, and frontend tracking server -# Used specifically for local development mode - -set -e # Exit on any error - -# Configuration -RF_MLFLOW_PORT=${RF_MLFLOW_PORT:=8852} -RF_MLFLOW_HOST=${RF_MLFLOW_HOST:=127.0.0.1} -RF_FRONTEND_PORT=${RF_FRONTEND_PORT:=8853} -RF_FRONTEND_HOST=${RF_FRONTEND_HOST:=0.0.0.0} -# API server configuration - these should match DispatcherConfig in constants.py -RF_API_PORT=${RF_API_PORT:=8851} -RF_API_HOST=${RF_API_HOST:=127.0.0.1} - -RF_DB_PATH="${RF_DB_PATH:=$HOME/db}" -# Directory paths -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" -# Navigate to rapidfireai/fit directory from setup/fit directory -RAPIDFIRE_FIT_DIR="$PROJECT_ROOT/rapidfireai/fit" -DISPATCHER_DIR="$RAPIDFIRE_FIT_DIR/dispatcher" -FRONTEND_DIR="$RAPIDFIRE_FIT_DIR/frontend" - -# Colors for output -RED='\033[0;31m' -GREEN='\033[0;32m' -YELLOW='\033[1;33m' -BLUE='\033[0;34m' -NC='\033[0m' # No Color - -# PID file to track processes -RF_PID_FILE="${RF_PID_FILE:=$HOME/rapidfire_pids.txt}" - -# Function to print colored output -print_status() { - echo -e "${BLUE}[$(date '+%Y-%m-%d %H:%M:%S')]${NC} $1" -} - -print_success() { - echo -e "${GREEN}[$(date '+%Y-%m-%d %H:%M:%S')]${NC} $1" -} - -print_error() { - echo -e "${RED}[$(date '+%Y-%m-%d %H:%M:%S')]${NC} $1" -} - -print_warning() { - echo -e "${YELLOW}[$(date '+%Y-%m-%d %H:%M:%S')]${NC} $1" -} - -# Function to setup Python environment -setup_python_env() { - print_status "Setting up Python environment..." - - # Check if we're in a virtual environment - # Check multiple indicators: VIRTUAL_ENV env var, conda environment, python path, and pip path - local in_venv=false - - if [[ -n "$VIRTUAL_ENV" ]]; then - in_venv=true - print_status "Detected virtual environment: $VIRTUAL_ENV" - elif [[ -n "$CONDA_DEFAULT_ENV" ]]; then - in_venv=true - print_status "Detected conda environment: $CONDA_DEFAULT_ENV" - elif [[ "$(which python 2>/dev/null)" == *"venv"* ]] || [[ "$(which python 2>/dev/null)" == *"virtualenv"* ]]; then - in_venv=true - print_status "Detected virtual environment via Python path" - elif [[ "$(which pip 2>/dev/null)" == *"venv"* ]] || [[ "$(which pip 2>/dev/null)" == *"virtualenv"* ]]; then - in_venv=true - print_status "Detected virtual environment via pip path" - elif [[ "$(which python 2>/dev/null)" == *"conda"* ]] || [[ "$(which python 2>/dev/null)" == *"miniconda"* ]]; then - in_venv=true - print_status "Detected conda environment via Python path" - elif [[ "$(which pip 2>/dev/null)" == *"conda"* ]] || [[ "$(which pip 2>/dev/null)" == *"miniconda"* ]]; then - in_venv=true - print_status "Detected conda environment via pip path" - fi - - if [[ "$in_venv" == "false" ]]; then - print_warning "Not in a virtual environment. This may cause permission issues." - print_status "Attempting to install with --user flag to avoid permission issues..." - PIP_USER_FLAG="--user" - else - PIP_USER_FLAG="" - fi - - # Debug information - print_status "Environment info:" - print_status " VIRTUAL_ENV: ${VIRTUAL_ENV:-'not set'}" - print_status " CONDA_DEFAULT_ENV: ${CONDA_DEFAULT_ENV:-'not set'}" - print_status " Python path: $(which python 2>/dev/null || echo 'not found')" - print_status " Pip path: $(which pip 2>/dev/null || echo 'not found')" - print_status " Using --user flag: $([[ -n "$PIP_USER_FLAG" ]] && echo 'yes' || echo 'no')" - - # Install rapidfireai in development mode - cd "$PROJECT_ROOT" - print_status "Installing rapidfireai in development mode..." - - if pip install -e . $PIP_USER_FLAG; then - print_success "rapidfireai package installed successfully" - else - print_error "Failed to install rapidfireai package" - print_warning "If you're getting permission errors, try:" - print_warning "1. Activate a virtual environment: python3 -m venv .venv && source .venv/bin/activate" - print_warning "2. Or run with sudo (not recommended): sudo pip install -e ." - return 1 - fi - - # Install requirements - if [[ -f "requirements.txt" ]]; then - print_status "Installing Python requirements..." - if pip install -r requirements.txt $PIP_USER_FLAG; then - print_success "Python requirements installed successfully" - else - print_error "Failed to install Python requirements" - return 1 - fi - fi - - cd "$SCRIPT_DIR" # Return to script directory - return 0 -} - -# Function to cleanup processes on exit -cleanup() { - print_warning "Shutting down services..." - - # Stop Docker containers if running - if docker ps --format "table {{.Names}}" | grep -q "rapidfire-frontend"; then - print_status "Stopping Docker container: rapidfire-frontend" - docker stop rapidfire-frontend >/dev/null 2>&1 || true - docker rm rapidfire-frontend >/dev/null 2>&1 || true - fi - - # Kill processes by port (more reliable for MLflow) - for port in $RF_MLFLOW_PORT $RF_FRONTEND_PORT $RF_API_PORT; do - local pids=$(lsof -ti :$port 2>/dev/null || true) - if [[ -n "$pids" ]]; then - print_status "Killing processes on port $port" - echo "$pids" | xargs kill -TERM 2>/dev/null || true - sleep 2 - # Force kill if still running - local remaining_pids=$(lsof -ti :$port 2>/dev/null || true) - if [[ -n "$remaining_pids" ]]; then - echo "$remaining_pids" | xargs kill -9 2>/dev/null || true - fi - fi - done - - # Clean up tracked PIDs - if [[ -f "$RF_PID_FILE" ]]; then - while read -r pid service; do - if kill -0 "$pid" 2>/dev/null; then - print_status "Stopping $service (PID: $pid)" - # Kill process group to get child processes too - kill -TERM -$pid 2>/dev/null || kill -TERM $pid 2>/dev/null || true - sleep 1 - # Force kill if still running - if kill -0 "$pid" 2>/dev/null; then - kill -9 -$pid 2>/dev/null || kill -9 $pid 2>/dev/null || true - fi - fi - done < "$RF_PID_FILE" - rm -f "$RF_PID_FILE" - fi - - # Final cleanup - kill any remaining MLflow or gunicorn processes - pkill -f "mlflow server" 2>/dev/null || true - pkill -f "gunicorn.*rapidfireai" 2>/dev/null || true - - print_success "All services stopped" - exit 0 -} - -# Function to check if a port is available -check_port() { - local port=$1 - local service=$2 - - if lsof -Pi :$port -sTCP:LISTEN -t >/dev/null 2>&1; then - print_error "Port $port is already in use. Cannot start $service." - return 1 - fi - return 0 -} - -# Function to wait for service to be ready -wait_for_service() { - local host=$1 - local port=$2 - local service=$3 - local max_attempts=${4:-30} # Allow custom timeout, default 30 seconds - local attempt=1 - - print_status "Waiting for $service to be ready on $host:$port (timeout: ${max_attempts}s)..." - - while [ $attempt -le $max_attempts ]; do - if nc -z "$host" "$port" 2>/dev/null; then - print_success "$service is ready!" - return 0 - fi - sleep 1 - ((attempt++)) - done - - print_error "$service failed to start within expected time (${max_attempts}s)" - return 1 -} - -# Function to start MLflow server -start_mlflow() { - print_status "Starting MLflow server..." - - if ! check_port $RF_MLFLOW_PORT "MLflow server"; then - return 1 - fi - - # Start MLflow server in background with process group - # Use setsid on Linux, nohup on macOS - if command -v setsid &> /dev/null; then - setsid mlflow server \ - --host $RF_MLFLOW_HOST \ - --port $RF_MLFLOW_PORT \ - --backend-store-uri sqlite:///${RF_DB_PATH}/rapidfire_mlflow.db > /dev/null 2>&1 & - else - nohup mlflow server \ - --host $RF_MLFLOW_HOST \ - --port $RF_MLFLOW_PORT \ - --backend-store-uri sqlite:///${RF_DB_PATH}/rapidfire_mlflow.db > /dev/null 2>&1 & - fi - - local mlflow_pid=$! - echo "$mlflow_pid MLflow" >> "$RF_PID_FILE" - - # Wait for MLflow to be ready - if wait_for_service $RF_MLFLOW_HOST $RF_MLFLOW_PORT "MLflow server"; then - print_success "MLflow server started (PID: $mlflow_pid)" - print_status "MLflow UI available at: http://$RF_MLFLOW_HOST:$RF_MLFLOW_PORT" - return 0 - else - return 1 - fi -} - -# Function to start API server -start_api_server() { - print_status "Starting API server with Gunicorn..." - - # Check if dispatcher directory exists - if [[ ! -d "$DISPATCHER_DIR" ]]; then - print_error "Dispatcher directory not found at $DISPATCHER_DIR" - return 1 - fi - - # Check if gunicorn config file exists - if [[ ! -f "$DISPATCHER_DIR/gunicorn.conf.py" ]]; then - print_error "gunicorn.conf.py not found in dispatcher directory" - return 1 - fi - - # Create database directory - print_status "Creating database directory..." - mkdir -p ~/db - # Ensure proper permissions - chmod 755 ~/db - - # Change to dispatcher directory and start Gunicorn server - cd "$DISPATCHER_DIR" - - # Set PYTHONPATH to include the project root - export PYTHONPATH="$PROJECT_ROOT:$PYTHONPATH" - - # Start Gunicorn server in background - gunicorn -c gunicorn.conf.py & - - local api_pid=$! - cd "$SCRIPT_DIR" # Return to original directory - echo "$api_pid API_Server" >> "$RF_PID_FILE" - - # Wait for API server to be ready - if wait_for_service $RF_API_HOST $RF_API_PORT "API server"; then - print_success "API server started (PID: $api_pid)" - print_status "API server available at: http://$RF_API_HOST:$RF_API_PORT" - return 0 - else - return 1 - fi -} - -# Function to build and start frontend server -start_frontend() { - print_status "Starting frontend tracking server..." - - if ! check_port $RF_FRONTEND_PORT "Frontend server"; then - return 1 - fi - - # Check if frontend directory exists - if [[ ! -d "$FRONTEND_DIR" ]]; then - print_error "Frontend directory not found at $FRONTEND_DIR" - return 1 - fi - - # Change to frontend directory - cd "$FRONTEND_DIR" - - # Check if we should use Node.js (preferred) or Docker - print_status "Starting frontend with Node.js directly..." - - # Determine which package manager to use - local yarn_binary="" - if [[ -f ".yarnrc.yml" ]]; then - # Check for available yarn releases in the yarn/releases directory - if [[ -f "yarn/releases/yarn-4.9.1.cjs" ]]; then - yarn_binary="yarn/releases/yarn-4.9.1.cjs" - elif [[ -f "yarn/releases/yarn-4.6.0.cjs" ]]; then - yarn_binary="yarn/releases/yarn-4.6.0.cjs" - elif [[ -f "yarn/releases/yarn-3.5.0.cjs" ]]; then - yarn_binary="yarn/releases/yarn-3.5.0.cjs" - fi - fi - - # Check if node_modules exists - if [[ ! -d "node_modules" ]]; then - print_status "Installing Node.js dependencies..." - # Check if this is a Yarn 2+ project (has .yarnrc.yml) - if [[ -f ".yarnrc.yml" ]]; then - print_status "Using local Yarn binary..." - - if [[ -n "$yarn_binary" ]]; then - node "$yarn_binary" install || { - print_error "Failed to install dependencies with local yarn ($yarn_binary)" - cd "$SCRIPT_DIR" - return 1 - } - else - print_error "No local yarn binary found in yarn/releases/" - cd "$SCRIPT_DIR" - return 1 - fi - elif command -v yarn &> /dev/null; then - yarn install || { - print_error "Failed to install dependencies with yarn" - cd "$SCRIPT_DIR" - return 1 - } - else - npm install || { - print_error "Failed to install dependencies with npm" - cd "$SCRIPT_DIR" - return 1 - } - fi - fi - - # Start Node.js server with appropriate package manager - print_status "Starting development server..." - print_status "Frontend logs will be written to: $SCRIPT_DIR/frontend.log" - - # Use yarn if available, otherwise fall back to npm - if [[ -f ".yarnrc.yml" ]] && [[ -n "$yarn_binary" ]]; then - print_status "Using local Yarn binary to start server..." - PORT=$RF_FRONTEND_PORT nohup node "$yarn_binary" start > "$SCRIPT_DIR/frontend.log" 2>&1 & - elif command -v yarn &> /dev/null; then - print_status "Using system Yarn to start server..." - PORT=$RF_FRONTEND_PORT nohup yarn start > "$SCRIPT_DIR/frontend.log" 2>&1 & - else - print_status "Using npm to start server..." - PORT=$RF_FRONTEND_PORT nohup npm start > "$SCRIPT_DIR/frontend.log" 2>&1 & - fi - - local frontend_pid=$! - cd "$SCRIPT_DIR" # Return to original directory - echo "$frontend_pid Frontend_Node" >> "$RF_PID_FILE" - - # Wait for frontend to be ready with longer timeout for development server - if wait_for_service localhost $RF_FRONTEND_PORT "Frontend server" 120; then - print_success "Frontend server started with Node.js (PID: $frontend_pid)" - print_status "Frontend available at: http://localhost:$RF_FRONTEND_PORT" - return 0 - else - print_error "Frontend development server failed to start. Showing recent logs:" - if [[ -f "$SCRIPT_DIR/frontend.log" ]]; then - echo "=== Last 20 lines of frontend.log ===" - tail -20 "$SCRIPT_DIR/frontend.log" - echo "=== End of logs ===" - else - print_error "No frontend.log file found" - fi - return 1 - fi -} - -# Function to display running services -show_status() { - print_status "RapidFire AI Services Status:" - echo "==================================" - - if [[ -f "$RF_PID_FILE" ]]; then - while read -r pid service; do - if kill -0 "$pid" 2>/dev/null; then - print_success "$service is running (PID: $pid)" - else - print_error "$service is not running (PID: $pid)" - fi - done < "$RF_PID_FILE" - else - print_warning "No services are currently tracked" - fi - - # Check Docker container - if docker ps -q -f name=rapidfire-frontend | grep -q .; then - print_success "Frontend Docker container is running" - fi - - echo "" - print_status "Available endpoints:" - echo "- MLflow UI: http://$RF_MLFLOW_HOST:$RF_MLFLOW_PORT" - echo "- Frontend: http://$RF_FRONTEND_HOST:$RF_FRONTEND_PORT" - echo "- API Server: http://$RF_API_HOST:$RF_API_PORT" -} - -# Main execution -main() { - print_status "Starting RapidFire AI services..." - - # Remove old PID file - rm -f "$RF_PID_FILE" - - # Set up signal handlers for cleanup - trap cleanup SIGINT SIGTERM EXIT - - # Check for required commands - for cmd in mlflow gunicorn; do - if ! command -v $cmd &> /dev/null; then - print_error "$cmd is not installed or not in PATH" - exit 1 - fi - done - - # Setup Python environment - if ! setup_python_env; then - print_error "Failed to setup Python environment" - exit 1 - fi - - # Start services - if start_mlflow && start_api_server && start_frontend; then - print_success "All services started successfully!" - show_status - - print_status "Press Ctrl+C to stop all services" - - # Keep script running and monitor processes - while true; do - sleep 5 - # Check if any process died - if [[ -f "$RF_PID_FILE" ]]; then - while read -r pid service; do - if ! kill -0 "$pid" 2>/dev/null; then - print_error "$service (PID: $pid) has stopped unexpectedly" - fi - done < "$RF_PID_FILE" - fi - done - else - print_error "Failed to start one or more services" - cleanup - exit 1 - fi -} - -# Handle command line arguments -case "${1:-start}" in - "start") - main - ;; - "stop") - cleanup - ;; - "status") - show_status - ;; - "restart") - cleanup - sleep 2 - main - ;; - "setup") - setup_python_env - ;; - *) - echo "Usage: $0 {start|stop|status|restart|setup}" - echo " start - Start all services (default)" - echo " stop - Stop all services" - echo " status - Show service status" - echo " restart - Restart all services" - echo " setup - Setup Python environment only" - exit 1 - ;; -esac diff --git a/setup/fit/trial.ipynb b/setup/fit/trial.ipynb deleted file mode 100644 index 5c37edfb..00000000 --- a/setup/fit/trial.ipynb +++ /dev/null @@ -1,86 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "498150c7", - "metadata": {}, - "source": [ - "### Barebones DPO (For setup-testing only)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "7a9251a4", - "metadata": {}, - "outputs": [], - "source": [ - "# Import required libraries for DPO (Direct Preference Optimization) training\n", - "from datasets import load_dataset\n", - "from trl import DPOConfig, DPOTrainer\n", - "from transformers import AutoModelForCausalLM, AutoTokenizer\n", - "\n", - "# Load the pre-trained Qwen model and tokenizer\n", - "model_name = \"Qwen/Qwen2-0.5B-Instruct\"\n", - "model = AutoModelForCausalLM.from_pretrained(model_name)\n", - "tokenizer = AutoTokenizer.from_pretrained(model_name)\n", - "\n", - "# Load and prepare datasets (using subset for faster training)\n", - "train_dataset = load_dataset(\n", - " \"trl-lib/ultrafeedback_binarized\", \n", - " split=\"train\"\n", - ").select(range(800)) # Use only first 1000 samples for demo\n", - "\n", - "eval_dataset = load_dataset(\n", - " \"trl-lib/ultrafeedback_binarized\", \n", - " split=\"test\"\n", - ").select(range(200)) # Use only first 1000 samples for demo\n", - "\n", - "# Configure DPO training parameters\n", - "training_args = DPOConfig(\n", - " output_dir=\"Qwen2-0.5B-DPO\", # Directory to save model checkpoints\n", - " per_device_train_batch_size=4, # Batch size per GPU/device\n", - " max_steps=50, # Total training steps\n", - " eval_strategy=\"steps\", # Evaluate\n", - " eval_steps=10, # Evaluate every 25 steps\n", - " report_to=\"tensorboard\" # Log metrics to TensorBoard\n", - ")\n", - "\n", - "# Initialize the DPO trainer\n", - "trainer = DPOTrainer(\n", - " model=model,\n", - " args=training_args,\n", - " processing_class=tokenizer,\n", - " train_dataset=train_dataset,\n", - " eval_dataset=eval_dataset\n", - ")\n", - "\n", - "# Start training\n", - "print(\"Starting DPO training...\")\n", - "trainer.train()\n", - "print(\"Training completed!\")" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python (base)", - "language": "python", - "name": "base" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.11" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/setup/evals/requirements-colab.txt b/setup/rapidfireai/requirements-colab.txt similarity index 60% rename from setup/evals/requirements-colab.txt rename to setup/rapidfireai/requirements-colab.txt index 3aac9c1b..bed3f35d 100644 --- a/setup/evals/requirements-colab.txt +++ b/setup/rapidfireai/requirements-colab.txt @@ -1,6 +1,9 @@ -# Colab Packages -# evals dependencies +# RapidFire AI - Google Colab Dependencies +# Supports both training (fit) and inference (evals) modes + # Core ML/AI Framework +torch>=2.8.0 +transformers>=4.56.1 sentence-transformers>=5.1.0 # Distributed Computing @@ -10,6 +13,15 @@ ray[default]>=2.51.0 # LLM Inference transformers>=4.56.1,<5.0.0 +# Fine-tuning (fit mode) +peft>=0.17.0 +trl>=0.21.0 +bitsandbytes>=0.47.0 +evaluate>=0.4.5 +rouge-score>=0.1.2 +nltk>=3.9.1 +sentencepiece>=0.2.1 + # OpenAI API openai>=1.106.1 tiktoken>=0.12.0 @@ -24,26 +36,43 @@ langchain-huggingface>=1.0.0 langchain-pinecone>=0.2.13 langchain-postgres>=0.0.17 -# Data Manipulation & Display -unstructured>=0.18.15 - -# Other -requests==2.32.5 +# LLM Inference (evals mode) vllm==0.11.0 -faiss-gpu-cu12 -fsspec>=2023.1.0,<=2025.3.0 # flashinfer excluded: v0.5.x requires compute capability >= 8.0 (sm_80+) # Colab T4 is sm_75; flashinfer's CUTLASS CuTE DSL also JIT-compiles at import # time inside vLLM's MP subprocess before CUDA is initialized, causing NVVM errors. # vLLM 0.11.0 falls back to built-in implementations without flashinfer. -# MLflow dashboard dependencies +# Data Manipulation & Display +unstructured>=0.18.15 + +# Vector Search +faiss-gpu-cu12 + +# Data Manipulation & Display +pandas>=2.0.0,<2.3.0 +fsspec>=2023.1.0,<=2025.3.0 + +# Notebook & Visualization +jupyter>=1.1.1 +ipywidgets>=7.3.4,<9.0.0 +tensorboard>=2.11.0 +ipython>=7.34.0 + +# Experiment Tracking & API mlflow>=3.2.0 gunicorn>=23.0.0 flask-cors>=5.0.1 +trackio # Logging loguru numpy==2.0.1 -protobuf<6.0.0 \ No newline at end of file +protobuf<6.0.0 + +# Utilities +psutil>=5.9.0 +requests>=2.32.0 +pytest>=8.4.1 +dill diff --git a/setup/rapidfireai/requirements-local.txt b/setup/rapidfireai/requirements-local.txt new file mode 100644 index 00000000..d057366a --- /dev/null +++ b/setup/rapidfireai/requirements-local.txt @@ -0,0 +1,67 @@ +# RapidFire AI - Local/Server Dependencies +# Supports both training (fit) and inference (evals) modes + +# Core ML/AI Framework +torch>=2.5.1 +torchvision>=0.20.1 +torchaudio>=2.5.1 +transformers>=4.56.1 +sentence-transformers>=5.1.0 +datasets>=3.6.0 +gpustat>=1.1.1 + +# Distributed Computing +ray>=2.49.0 +ray[default]>=2.49.0 + +# Fine-tuning (fit mode) +peft>=0.17.0 +trl>=0.21.0 +bitsandbytes>=0.47.0 +evaluate>=0.4.5 +rouge-score>=0.1.2 +nltk>=3.9.1 +sentencepiece>=0.2.1 + +# OpenAI API +openai>=1.106.1 +tiktoken>=0.12.0 + +# LangChain Ecosystem +langchain>=1.0.5 +langchain-classic>=1.0.0 +langchain-core>=1.0.4 +langchain-community>=0.4.1 +langchain-openai>=1.0.2 +langchain-huggingface>=1.0.0 + +# Vector Search +faiss-gpu-cu12>=1.13.0 + +# Data Manipulation & Display +pandas>=2.3.2 +pyarrow>=21.0.0 +numpy>=1.26.4 +scipy>=1.16.1 +unstructured>=0.18.15 + +# Notebook & Visualization +jupyter>=1.1.1 +ipykernel>=6.30.1 +ipywidgets>=7.3.4,<9.0.0 +tensorboard>=2.11.0 +ipython>=7.34.0 + +# Experiment Tracking & API +mlflow>=3.2.0 +gunicorn>=23.0.0 +flask-cors>=5.0.1 +trackio + +# Utilities +psutil>=7.0.0 +tqdm>=4.67.1 +typing-extensions>=4.0.0 +requests>=2.32.0 +grpcio +dill diff --git a/setup/rapidfireai/requirements.txt b/setup/rapidfireai/requirements.txt new file mode 100644 index 00000000..5fa34032 --- /dev/null +++ b/setup/rapidfireai/requirements.txt @@ -0,0 +1,46 @@ +# RapidFire AI - Base Dependencies +# Used for pip install without environment-specific packages + +# Core ML/AI Framework +torch>=2.5.1 +transformers>=4.56.1 +sentence-transformers>=5.1.0 +datasets>=3.6.0 + +# Distributed Computing +ray>=2.49.0 + +# Fine-tuning +peft>=0.17.0 +trl>=0.21.0 +bitsandbytes>=0.47.0 +evaluate>=0.4.5 + +# OpenAI API +openai>=1.106.1 +tiktoken>=0.12.0 + +# LangChain Ecosystem +langchain>=1.0.5 +langchain-core>=1.0.4 +langchain-community>=0.4.1 +langchain-openai>=1.0.2 +langchain-huggingface>=1.0.0 + +# Data Manipulation +pandas>=2.0.0 +numpy>=1.26.4 + +# Notebook +jupyter>=1.1.1 +ipywidgets>=7.3.4,<9.0.0 + +# Experiment Tracking & API +mlflow>=3.2.0 +gunicorn>=23.0.0 +flask-cors>=5.0.1 + +# Utilities +psutil>=5.9.0 +requests>=2.32.0 +dill diff --git a/setup/start.sh b/setup/start.sh index 7fdeecea..d605480c 100755 --- a/setup/start.sh +++ b/setup/start.sh @@ -70,7 +70,7 @@ RF_PID_FILE="${RF_PID_FILE:=$RF_HOME/rapidfire_pids.txt}" # Directory paths for pip-installed package SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -# Navigate to rapidfireai/fit directory from setup/fit directory +# Navigate to rapidfireai package directory from setup directory RAPIDFIRE_DIR="$SCRIPT_DIR/../rapidfireai" RAPIDFIRE_FIT_DIR="$RAPIDFIRE_DIR/fit" RAPIDFIRE_EVALS_DIR="$RAPIDFIRE_DIR/evals" @@ -80,7 +80,7 @@ else FRONTEND_DIR="$RAPIDFIRE_DIR/frontend" fi RAPIDFIRE_MODE=$(cat $RF_HOME/rf_mode.txt 2>/dev/null || echo "fit") -DISPATCHER_DIR="$RAPIDFIRE_DIR/$RAPIDFIRE_MODE/dispatcher" +DISPATCHER_DIR="$RAPIDFIRE_DIR/dispatcher" RF_PYTHON_EXECUTABLE=${RF_PYTHON_EXECUTABLE:-python3} RF_PIP_EXECUTABLE=${RF_PIP_EXECUTABLE:-pip3} @@ -194,12 +194,15 @@ cleanup() { rm -f "$RF_PID_FILE" fi + # Stop Ray cluster + stop_ray + # Final cleanup - ONLY if NOT in Colab mode # Colab mode skips this to avoid killing Jupyter/IPython infrastructure if [[ "$RF_COLAB_MODE" != "true" ]]; then # Safe, specific patterns for non-Colab environments pkill -f "mlflow server" 2>/dev/null || true - pkill -f "gunicorn.*rapidfireai.$RAPIDFIRE_MODE.dispatcher" 2>/dev/null || true + pkill -f "gunicorn.*rapidfireai.dispatcher" 2>/dev/null || true # Only kill Flask server if we're not in Colab (frontend doesn't run in Colab) pkill -f "python.*rapidfireai/frontend/server.py" 2>/dev/null || true # Stop Converge if it was running @@ -231,7 +234,7 @@ ping_port() { if command -v nc &> /dev/null; then ping_command="$(command -v nc) -z $host $port" else - ping_command="$RF_PYTHON_EXECUTABLE -c 'from rapidfireai.utils.ping import ping_server; checker=ping_server(\"${host}\", ${port}); exit(1) if not checker else exit(0)'" + ping_command="$RF_PYTHON_EXECUTABLE -c 'from rapidfireai.platform.ping import ping_server; checker=ping_server(\"${host}\", ${port}); exit(1) if not checker else exit(0)'" fi eval $ping_command return $? @@ -322,7 +325,7 @@ wait_for_service() { if command -v nc &> /dev/null; then ping_command="$(command -v nc) -z $host $port" else - ping_command="$RF_PYTHON_EXECUTABLE -c 'from rapidfireai.utils.ping import ping_server; checker=ping_server(\"${host}\", ${port}); exit(1) if not checker else exit(0)'" + ping_command="$RF_PYTHON_EXECUTABLE -c 'from rapidfireai.platform.ping import ping_server; checker=ping_server(\"${host}\", ${port}); exit(1) if not checker else exit(0)'" fi while [ $attempt -le $max_attempts ]; do if eval ${ping_command} &>/dev/null; then @@ -337,6 +340,47 @@ wait_for_service() { return 1 } +# Function to start Ray head node +start_ray() { + print_status "Starting Ray head node..." + + # Check if ray is installed + if ! command -v ray &> /dev/null; then + print_error "ray is not installed or not in PATH" + return 1 + fi + + # Check if Ray is already running + if ray status &>/dev/null 2>&1; then + print_success "Ray cluster already running, skipping start" + return 0 + fi + + # Start Ray head node with dashboard + ray start --head \ + --dashboard-host="$RF_RAY_HOST" \ + --dashboard-port="$RF_RAY_PORT" \ + --disable-usage-stats \ + > "$RF_LOG_PATH/ray.log" 2>&1 + + if [ $? -eq 0 ]; then + print_success "Ray head node started (dashboard at http://$RF_RAY_HOST:$RF_RAY_PORT)" + return 0 + else + print_error "Failed to start Ray head node. Check $RF_LOG_PATH/ray.log" + return 1 + fi +} + +# Function to stop Ray cluster +stop_ray() { + if ray status &>/dev/null 2>&1; then + print_status "Stopping Ray cluster..." + ray stop --force 2>/dev/null || true + print_success "Ray cluster stopped" + fi +} + # Function to start MLflow server start_mlflow() { print_status "Starting MLflow server..." @@ -744,19 +788,24 @@ show_status() { else print_error "๐Ÿšจ RapidFire API server is not ready!" fi - # if [[ "$rf_mode" == "evals" ]]; then - # if ping_port $RF_RAY_HOST $RF_RAY_PORT; then - # print_success "๐Ÿš€ RapidFire Ray server is ready!" - # else - # print_error "๐Ÿšจ RapidFire Ray server is not ready!" - # fi - # fi + if ping_port $RF_RAY_HOST $RF_RAY_PORT; then + print_success "๐Ÿš€ Ray cluster is ready! (dashboard at http://$RF_RAY_HOST:$RF_RAY_PORT)" + else + print_error "๐Ÿšจ Ray cluster is not running!" + fi # Show log file status echo "" print_status "Log files:" - # Always check api.log + # Always check ray.log and api.log + if [[ -f "$RF_LOG_PATH/ray.log" ]]; then + local size=$(du -h "$RF_LOG_PATH/ray.log" | cut -f1) + print_status "- $RF_LOG_PATH/ray.log: $size" + else + print_warning "- $RF_LOG_PATH/ray.log: not found" + fi + if [[ -f "$RF_LOG_PATH/api.log" ]]; then local size=$(du -h "$RF_LOG_PATH/api.log" | cut -f1) print_status "- $RF_LOG_PATH/api.log: $size" @@ -788,7 +837,7 @@ show_status() { # Function to start services based on mode start_services() { local services_started=0 - local total_services=1 # API server always runs + local total_services=2 # Ray + API server always run # Calculate total services based on mode # MLflow runs unless tensorboard-only in Colab @@ -805,6 +854,13 @@ start_services() { print_status "Starting $total_services service(s)..." + # Start Ray head node (always) + if start_ray; then + ((services_started++)) + else + print_error "Failed to start Ray head node" + fi + # Start MLflow server (conditionally) if [[ "$RF_MLFLOW_ENABLED" == "true" ]]; then if start_mlflow; then @@ -869,7 +925,7 @@ main() { trap cleanup SIGINT SIGTERM EXIT # Check for required commands - for cmd in mlflow gunicorn; do + for cmd in ray mlflow gunicorn; do if ! command -v $cmd &> /dev/null; then print_error "$cmd is not installed or not in PATH" exit 1 diff --git a/tests/conftest.py b/tests/conftest.py index 8fe67e34..4ad15db5 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -31,7 +31,7 @@ def mock_mlflow_manager(): Returns: Mock: Mocked MLflowManager with all required methods """ - from rapidfireai.utils.metric_mlflow_manager import MLflowMetricLogger + from rapidfireai.metrics.metric_mlflow_manager import MLflowMetricLogger mock = Mock(spec=MLflowMetricLogger) mock.create_run.return_value = "test_run_id" @@ -76,7 +76,7 @@ def mlflow_logger(mock_mlflow_manager): Returns: MLflowMetricLogger: Logger with mocked backend """ - from rapidfireai.fit.utils.metric_logger import MLflowMetricLogger + from rapidfireai.metrics.metric_mlflow_manager import MLflowMetricLogger logger = MLflowMetricLogger("http://localhost:8852") logger.mlflow_manager = mock_mlflow_manager @@ -94,26 +94,11 @@ def tensorboard_logger(temp_tensorboard_dir): """ import os os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" - from rapidfireai.fit.utils.metric_logger import TensorBoardMetricLogger + from rapidfireai.metrics.metric_tensorboard_manager import TensorBoardMetricLogger return TensorBoardMetricLogger(temp_tensorboard_dir) -@pytest.fixture -def dual_logger(mock_mlflow_manager, temp_tensorboard_dir): - """ - Create a DualMetricLogger with mocked MLflow and temp TensorBoard. - - Returns: - DualMetricLogger: Logger with both backends - """ - from rapidfireai.fit.utils.metric_logger import DualMetricLogger - - logger = DualMetricLogger("http://localhost:8852", temp_tensorboard_dir) - logger.mlflow_logger.mlflow_manager = mock_mlflow_manager - - return logger - @pytest.fixture def sample_metrics(): diff --git a/tests/test_metric_logger.py b/tests/test_metric_logger.py index ce3f2247..03783238 100644 --- a/tests/test_metric_logger.py +++ b/tests/test_metric_logger.py @@ -1,11 +1,9 @@ """ Comprehensive test suite for metric logging abstraction layer. -Tests the MetricLogger interface and all implementations: +Tests the MetricLogger interface and implementations: - MLflowMetricLogger - TensorBoardMetricLogger -- DualMetricLogger -- create_metric_logger() factory function """ import os @@ -14,90 +12,9 @@ import pytest -from rapidfireai.fit.utils.metric_logger import ( - MetricLogger, - MLflowMetricLogger, - TensorBoardMetricLogger, - DualMetricLogger, - create_metric_logger, -) - - -class TestMetricLoggerFactory: - """Test suite for create_metric_logger() factory function.""" - - def test_create_mlflow_logger(self): - """Test creating MLflow logger with valid URI.""" - logger = create_metric_logger( - backend="mlflow", - mlflow_tracking_uri="http://localhost:8852" - ) - - assert isinstance(logger, MLflowMetricLogger) - - def test_create_tensorboard_logger(self, temp_tensorboard_dir): - """Test creating TensorBoard logger with valid log_dir.""" - logger = create_metric_logger( - backend="tensorboard", - tensorboard_log_dir=temp_tensorboard_dir - ) - - assert isinstance(logger, TensorBoardMetricLogger) - - def test_create_dual_logger(self, temp_tensorboard_dir): - """Test creating Dual logger with both parameters.""" - logger = create_metric_logger( - backend="both", - mlflow_tracking_uri="http://localhost:8852", - tensorboard_log_dir=temp_tensorboard_dir - ) - - assert isinstance(logger, DualMetricLogger) - - def test_backend_case_insensitive(self): - """Test that backend name is case-insensitive.""" - logger1 = create_metric_logger( - backend="MLFLOW", - mlflow_tracking_uri="http://localhost:8852" - ) - logger2 = create_metric_logger( - backend="MlFlow", - mlflow_tracking_uri="http://localhost:8852" - ) - - assert isinstance(logger1, MLflowMetricLogger) - assert isinstance(logger2, MLflowMetricLogger) - - def test_invalid_backend_name(self): - """Test error handling for invalid backend name.""" - with pytest.raises(ValueError, match="Invalid backend"): - create_metric_logger( - backend="invalid_backend", - mlflow_tracking_uri="http://localhost:8852" - ) - - def test_missing_mlflow_uri(self): - """Test error when mlflow_tracking_uri is missing for mlflow backend.""" - with pytest.raises(ValueError, match="mlflow_tracking_uri required"): - create_metric_logger(backend="mlflow") - - def test_missing_tensorboard_dir(self): - """Test error when tensorboard_log_dir is missing for tensorboard backend.""" - with pytest.raises(ValueError, match="tensorboard_log_dir required"): - create_metric_logger(backend="tensorboard") - - def test_missing_both_parameters_for_dual(self): - """Test error when both parameters are missing for 'both' backend.""" - with pytest.raises(ValueError, match="Both mlflow_tracking_uri and tensorboard_log_dir required"): - create_metric_logger(backend="both") - - def test_missing_mlflow_uri_for_dual(self, temp_tensorboard_dir): - """Test error when only TensorBoard dir provided for 'both' backend.""" - with pytest.raises(ValueError, match="Both mlflow_tracking_uri and tensorboard_log_dir required"): - create_metric_logger( - backend="both", - tensorboard_log_dir=temp_tensorboard_dir - ) +from rapidfireai.metrics.metric_logger import MetricLogger +from rapidfireai.metrics.metric_mlflow_manager import MLflowMetricLogger +from rapidfireai.metrics.metric_tensorboard_manager import TensorBoardMetricLogger class TestMLflowMetricLogger: @@ -277,87 +194,6 @@ def test_cleanup_on_delete(self, tensorboard_logger): mock_close2.assert_called_once() -class TestDualMetricLogger: - """Test suite for DualMetricLogger.""" - - def test_create_run_calls_both_backends(self, dual_logger, mock_mlflow_manager): - """Test that create_run calls both MLflow and TensorBoard.""" - with patch.object(dual_logger.tensorboard_logger, 'create_run', return_value="run_1") as mock_tb: - run_id = dual_logger.create_run("run_1") - - mock_mlflow_manager.create_run.assert_called_once_with("run_1") - mock_tb.assert_called_once_with("run_1") - assert run_id == "test_run_id" # Returns MLflow run_id - - def test_log_param_calls_both_backends(self, dual_logger, mock_mlflow_manager): - """Test that log_param calls both MLflow and TensorBoard.""" - with patch.object(dual_logger.tensorboard_logger, 'log_param') as mock_tb: - dual_logger.log_param("run_1", "learning_rate", "1e-3") - - mock_mlflow_manager.log_param.assert_called_once_with("run_1", "learning_rate", "1e-3") - mock_tb.assert_called_once_with("run_1", "learning_rate", "1e-3") - - def test_log_metric_calls_both_backends(self, dual_logger, mock_mlflow_manager): - """Test that log_metric calls both MLflow and TensorBoard with same parameters.""" - with patch.object(dual_logger.tensorboard_logger, 'log_metric') as mock_tb: - dual_logger.log_metric("run_1", "loss", 0.5, step=100) - - mock_mlflow_manager.log_metric.assert_called_once_with("run_1", "loss", 0.5, step=100) - mock_tb.assert_called_once_with("run_1", "loss", 0.5, step=100) - - def test_end_run_calls_both_backends(self, dual_logger, mock_mlflow_manager): - """Test that end_run calls both MLflow and TensorBoard.""" - with patch.object(dual_logger.tensorboard_logger, 'end_run') as mock_tb: - dual_logger.end_run("run_1") - - mock_mlflow_manager.end_run.assert_called_once_with("run_1") - mock_tb.assert_called_once_with("run_1") - - def test_delete_run_only_calls_mlflow(self, dual_logger, mock_mlflow_manager): - """Test that delete_run only calls MLflow (TensorBoard doesn't support delete).""" - with patch.object(dual_logger.tensorboard_logger, 'delete_run') as mock_tb: - dual_logger.delete_run("run_1") - - mock_mlflow_manager.delete_run.assert_called_once_with("run_1") - # TensorBoard delete_run is not called - mock_tb.assert_not_called() - - def test_get_run_metrics_returns_mlflow_metrics(self, dual_logger, mock_mlflow_manager): - """Test that get_run_metrics returns MLflow metrics (primary source).""" - metrics = dual_logger.get_run_metrics("run_1") - - mock_mlflow_manager.get_run_metrics.assert_called_once_with("run_1") - assert "loss" in metrics - assert "accuracy" in metrics - - def test_get_experiment_returns_mlflow_experiment(self, dual_logger, mock_mlflow_manager): - """Test that get_experiment returns MLflow experiment ID.""" - exp_id = dual_logger.get_experiment("test_experiment") - - mock_mlflow_manager.get_experiment.assert_called_once_with("test_experiment") - assert exp_id == "test_experiment_id" - - def test_both_backends_receive_identical_data(self, dual_logger, mock_mlflow_manager): - """Test that both backends receive identical data for all operations.""" - with patch.object(dual_logger.tensorboard_logger, 'create_run') as mock_tb_create: - with patch.object(dual_logger.tensorboard_logger, 'log_param') as mock_tb_param: - with patch.object(dual_logger.tensorboard_logger, 'log_metric') as mock_tb_metric: - # Create run - dual_logger.create_run("run_1") - mock_mlflow_manager.create_run.assert_called_with("run_1") - mock_tb_create.assert_called_with("run_1") - - # Log param - dual_logger.log_param("run_1", "batch_size", "32") - mock_mlflow_manager.log_param.assert_called_with("run_1", "batch_size", "32") - mock_tb_param.assert_called_with("run_1", "batch_size", "32") - - # Log metric - dual_logger.log_metric("run_1", "loss", 0.5, step=100) - mock_mlflow_manager.log_metric.assert_called_with("run_1", "loss", 0.5, step=100) - mock_tb_metric.assert_called_with("run_1", "loss", 0.5, step=100) - - class TestMetricLoggerInterface: """Test suite for MetricLogger abstract base class.""" @@ -376,7 +212,7 @@ def test_all_implementations_have_required_methods(self): 'get_run_metrics', ] - for implementation in [MLflowMetricLogger, TensorBoardMetricLogger, DualMetricLogger]: + for implementation in [MLflowMetricLogger, TensorBoardMetricLogger]: for method in required_methods: assert hasattr(implementation, method) assert callable(getattr(implementation, method)) @@ -391,18 +227,12 @@ def test_tensorboard_logger_conforms_to_interface(self, temp_tensorboard_dir): logger = TensorBoardMetricLogger(temp_tensorboard_dir) assert isinstance(logger, MetricLogger) - def test_dual_logger_conforms_to_interface(self, temp_tensorboard_dir): - """Test that DualMetricLogger conforms to MetricLogger interface.""" - logger = DualMetricLogger("http://localhost:8852", temp_tensorboard_dir) - assert isinstance(logger, MetricLogger) - - class TestCallbacksIntegration: """Test suite for integration with callbacks.""" def test_metric_logging_callback_accepts_metric_logger(self): """Test that MetricLoggingCallback accepts metric_logger parameter.""" - from rapidfireai.ml.callbacks import MetricLoggingCallback + from rapidfireai.fit.ml.callbacks import MetricLoggingCallback mock_logger = Mock(spec=MetricLogger) callback = MetricLoggingCallback( @@ -417,7 +247,7 @@ def test_metric_logging_callback_accepts_metric_logger(self): def test_generation_metrics_callback_accepts_metric_logger(self): """Test that GenerationMetricsCallback accepts metric_logger parameter.""" - from rapidfireai.ml.callbacks import GenerationMetricsCallback + from rapidfireai.fit.ml.callbacks import GenerationMetricsCallback from transformers import AutoTokenizer from datasets import Dataset @@ -438,7 +268,7 @@ def test_generation_metrics_callback_accepts_metric_logger(self): def test_callback_calls_log_metric(self): """Test that callback calls metric_logger.log_metric().""" - from rapidfireai.ml.callbacks import MetricLoggingCallback + from rapidfireai.fit.ml.callbacks import MetricLoggingCallback from transformers import TrainerState, TrainerControl, TrainingArguments mock_logger = Mock(spec=MetricLogger) @@ -463,7 +293,7 @@ def test_callback_calls_log_metric(self): def test_callback_step_offset_works_correctly(self): """Test that callbacks apply completed_steps offset correctly.""" - from rapidfireai.ml.callbacks import MetricLoggingCallback + from rapidfireai.fit.ml.callbacks import MetricLoggingCallback from transformers import TrainerState, TrainerControl, TrainingArguments mock_logger = Mock(spec=MetricLogger)