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

rapidfireai/utils/worker_manager.py

@@ -0,0 +1,219 @@
+"""
+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 typing import Optional
+
+from torch.multiprocessing import Event, Process
+
+from rapidfireai.backend.worker import Worker
+from rapidfireai.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.
+    """
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(worker_id)
+
+    # 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: Optional[int] = None
+        self.processes: list[Process] = []
+        self.shutdown_events: list[EventType] = []
+        self.shutdown_event = threading.Event()
+        self.monitor_thread: Optional[threading.Thread] = 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
+                os.getpgid(self.parent_pid)
+                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
+        os.setpgrp()
+        self.process_group_id = os.getpgrp()
+
+        self.logger.debug(f"Starting worker processes in process group {self.process_group_id}")
+
+        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

rapidfireai/version.py

@@ -0,0 +1,6 @@
+"""
+Version information for RapidFire AI
+"""
+
+__version__ = "0.9.9"
+__version_info__ = (0, 9, 9) 

rapidfireai-0.9.9.dist-info/METADATA

@@ -0,0 +1,242 @@
+Metadata-Version: 2.4
+Name: rapidfireai
+Version: 0.9.9
+Summary: RapidFire AI - Machine Learning Platform
+Author-email: "RapidFire AI Inc." <support@rapidfire.ai>
+License: Apache-2.0
+Project-URL: Homepage, https://rapidfire.ai
+Keywords: ai,rapidfire,deep-learning,artificial-intelligence,machine-learning,mlflow,experiment-tracking
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Topic :: Software Development :: Libraries :: Application Frameworks
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: pandas>=2.3.1
+Requires-Dist: torch>=2.8.0
+Requires-Dist: transformers>=4.55.2
+Requires-Dist: peft>=0.17.0
+Requires-Dist: trl==0.21.0
+Requires-Dist: bitsandbytes>=0.47.0
+Requires-Dist: nltk>=3.9.1
+Requires-Dist: evaluate>=0.4.5
+Requires-Dist: rouge-score>=0.1.2
+Requires-Dist: sentencepiece>=0.2.1
+Requires-Dist: dill>=0.3.8
+Requires-Dist: mlflow>=3.2.0
+Requires-Dist: pytest>=8.4.1
+Requires-Dist: flask>=3.1.1
+Requires-Dist: flask-cors>=6.0.1
+Requires-Dist: requests>=2.32.5
+Requires-Dist: loguru>=0.7.3
+Requires-Dist: ipython>=8.37.0
+Requires-Dist: jupyter>=1.1.1
+Requires-Dist: ipywidgets>=8.1.7
+Requires-Dist: uv>=0.8.14
+Provides-Extra: dev
+Requires-Dist: pytest>=6.0; extra == "dev"
+Requires-Dist: pytest-cov>=2.0; extra == "dev"
+Requires-Dist: black>=21.0; extra == "dev"
+Requires-Dist: flake8>=3.8; extra == "dev"
+Requires-Dist: mypy>=0.800; extra == "dev"
+Requires-Dist: pre-commit>=2.0; extra == "dev"
+Provides-Extra: full
+Requires-Dist: torch>=1.9.0; extra == "full"
+Requires-Dist: transformers>=4.20.0; extra == "full"
+Requires-Dist: peft>=0.3.0; extra == "full"
+Requires-Dist: trl==0.21.0; extra == "full"
+Requires-Dist: bitsandbytes>=0.39.0; extra == "full"
+Dynamic: license-file
+
+# RapidFire AI
+
+Rapid experimentation for easier, faster, and more impactful fine-tuning and post-training for LLMs and other DL models.
+
+## Overview
+
+RapidFire AI is a new experiment execution framework that transforms your LLM customization experimentation from slow, sequential processes into rapid, intelligent workflows with hyperparallelized training, dynamic real-time experiment control, and automatic multi-GPU system orchestration.
+
+![Usage workflow of RapidFire AI](./usage.png)
+
+
+## Getting Started
+### Prerequisites
+- [NVIDIA GPU using the 7.x or 8.x Compute Capability](https://developer.nvidia.com/cuda-gpus)
+- [NVIDIA CUDA Toolkit 11.8+](https://developer.nvidia.com/cuda-toolkit-archive)
+- [Python 3.12.x](https://www.python.org/downloads/)
+- [PyTorch 2.7.1+](https://pytorch.org/get-started/previous-versions/) with corresponding forward compatible prebuilt CUDA binaries
+
+### Installation/Starting
+```bash
+virtualenv -p python3 oss_venv
+source oss_venv/bin/activate
+
+# from pypi
+pip install rapidfireai
+
+# install specific dependencies and initialize rapidfire
+rapidfire init
+
+# start the rapidfire server
+rapidfire start
+
+# open up example notebook and start experiment
+```
+
+
+
+### Troubleshooting
+
+If you encounter port conflicts, you can kill existing processes:
+```bash
+lsof -t -i:5002 | xargs kill -9  # mlflow
+lsof -t -i:8080 | xargs kill -9  # dispatcher
+lsof -t -i:3000 | xargs kill -9  # frontend server
+```
+
+## Documentation
+
+Browse or reference the full documentation, example use case tutorials, all API details, dashboard details, and more [here](https://rapidfire-ai-oss-docs.readthedocs-hosted.com/).
+
+
+## Key Features
+
+### MLflow Integration
+Full MLflow support for experiment tracking and metrics visualization. A named RapidFire AI experiment corresponds to an MLflow experiment for comprehensive governance
+
+### Interactive Control Operations (IC Ops)
+First-of-its-kind dynamic real-time control over runs in flight. Can be invoked through the dashboard:
+- Stop active runs; puts them in a dormant state
+- Resume stopped runs; makes them active again
+- Clone and modify existing runs, with or without warm starting from parent’s weights
+- Delete unwanted or failed runs
+
+### Multi-GPU Support
+The Scheduler automatically handles multiple GPUs on the machine and divides resources across all running configs for optimal resource utilization.
+
+### Search and AutoML Support
+Built-in procedures for searching over configuration knob combinations, including Grid Search and Random Search. Easy to integrate with AutoML procedures. Native support for some popular AutoML procedures and customized automation of IC Ops coming soon.
+
+
+## Directory Structure
+
+```
+rapidfireai/
+├── automl/          # Search and AutoML algorithms for knob tuning
+├── backend/         # Core backend components (controller, scheduler, worker)
+├── db/              # Database interface and SQLite operations
+├── dispatcher/      # Flask-based web API for UI communication
+├── frontend/         # Frontend components (dashboard, IC Ops implementation)
+├── ml/              # ML training utilities and trainer classes
+├── utils/           # Utility functions and helper modules
+└── experiment.py    # Main experiment lifecycle management
+```
+
+## Architecture
+
+RapidFire AI adopts a microservices-inspired loosely coupled distributed architecture with:
+- **Dispatcher**: Web API layer for UI communication
+- **Database**: SQLite for state persistence
+- **Controller**: Central orchestrator running in user process
+- **Workers**: GPU-based training processes
+- **Dashboard**: Experiment tracking and visualization dashboard
+
+This design enables efficient resource utilization while providing a seamless user experience for AI experimentation.
+
+
+## Components
+
+### Dispatcher
+The dispatcher provides a REST API interface for the web UI. It can be run via Flask as a single app or via Gunicorn to have it load balanced. Handles interactive control features and displays the current state of the runs in the experiment.
+
+### Database
+Uses SQLite for persistent storage of metadata of experiments, runs, and artifacts. The Controller also uses it to talk with Workers on scheduling state. A clean asynchronous interface for all DB operations, including experiment lifecycle management and run tracking.
+
+### Controller
+Runs as part of the user’s console or Notebook process. Orchestrates the entire training lifecycle including model creation, worker management, and scheduling. The `run_fit` logic handles sample preprocessing, model creation for given knob configurations, worker initialization, and continuous monitoring of training progress across distributed workers.
+
+### Worker
+Handles the actual model training and inference on the GPUs. Workers poll the Database for tasks, load dataset chunks, and execute training runs with checkpointing and progress reporting. Currently expects any given model for given batch size to fit on a single GPU.
+
+### Experiment
+Manages the complete experiment lifecycle, including creation, naming conventions, and cleanup. Experiments are automatically named with unique suffixes if conflicts exist, and all experiment metadata is tracked in the Database. An experiment's running tasks are automatically cancelled when the process ends abruptly.
+
+### Dashboard
+A fork of MLflow that enables full tracking and visualization of all experiments and runs. It features a new panel for Interactive Control Ops that can be performed on any active runs.
+
+
+## Developing with RapidFire AI
+### Prerequisites
+- Python 3.x
+- Git
+- Ubuntu/Debian system (for apt package manager)
+
+```bash
+# Run these commands one after the other on a fresh Ubuntu machine
+
+# install dependencies
+sudo apt update -y
+
+# clone the repository
+git clone https://github.com/RapidFireAI/rapidfireai.git
+
+# navigate to the repository
+cd ./rapidfireai
+
+# install basic dependencies
+sudo apt install -y python3-virtualenv
+virtualenv -p python3 oss_venv
+source oss_venv/bin/activate
+pip3 install ipykernel
+pip3 install jupyter
+pip3 install "huggingface-hub[cli]"
+export PATH="$HOME/.local/bin:$PATH"
+hf auth login --token <your_token>
+
+# checkout the develop branch
+git checkout develop
+
+# install the repository as a python package
+pip3 install -r requirements.txt
+
+# Install correct version of vllm and flash-attn
+# uv pip install vllm=0.10.1.1 --torch-backend=cu126 or cu118
+# uv pip install flash-attn==1.0.9 --no-build-isoloation or 2.8.3
+
+# install frontend packages
+curl -fsSL https://deb.nodesource.com/setup_22.x | sudo -E bash - && sudo apt-get install -y nodejs
+npm install npm@10.5.1
+sudo apt install -y yarn
+
+# if running into node versioning errors, remove the previous version of node then run the lines above again
+sudo apt-get remove --purge nodejs libnode-dev libnode72 npm
+sudo apt autoremove --purge
+
+# check installations
+node -v # 22.x
+npm -v # 10.5.1
+
+# still inside venv, run the start script to begin all 3 servers
+chmod +x ./rapidfireai/start_dev.sh
+./rapidfireai/start_dev.sh start
+
+# run the notebook from within your IDE
+# make sure the notebook is running in the oss_venv virtual environment
+# head to settings in Cursor/VSCode and search for venv and add the path - $HOME/rapidfireai/oss_venv
+# we cannot run a Jupyter notebook directly since there are restrictions on Jupyter being able to create child processes
+
+# VSCode can port-forward localhost:3000 where the rf-frontend server will be running
+
+# for port clash issues -
+lsof -t -i:8080 | xargs kill -9 # dispatcher
+lsof -t -i:5002 | xargs kill -9 # mlflow
+lsof -t -i:3000 | xargs kill -9 # frontend
+```
+