rapidfireai 0.10.2rc4__tar.gz → 0.10.3rc1__tar.gz

{rapidfireai-0.10.2rc4/rapidfireai.egg-info → rapidfireai-0.10.3rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rapidfireai
-Version: 0.10.2rc4
+Version: 0.10.3rc1
 Summary: RapidFire AI: Rapid Experimentation Engine for Customizing LLMs
 Author-email: "RapidFire AI Inc." <support@rapidfire.ai>
 License: Apache-2.0
@@ -22,6 +22,7 @@ Requires-Python: >=3.12
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: pandas>=2.3.1
+Requires-Dist: psutil>=5.9.0
 Requires-Dist: torch>=2.8.0
 Requires-Dist: transformers>=4.55.2
 Requires-Dist: peft>=0.17.0
@@ -40,11 +41,13 @@ 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: ipywidgets<9.0.0,>=7.3.4
 Requires-Dist: uv>=0.8.14
+Requires-Dist: tensorboard>=2.11.0
 Provides-Extra: dev
 Requires-Dist: pytest>=6.0; extra == "dev"
 Requires-Dist: pytest-cov>=2.0; extra == "dev"
+Requires-Dist: pytest-mock>=3.10.0; extra == "dev"
 Requires-Dist: black>=21.0; extra == "dev"
 Requires-Dist: flake8>=3.8; extra == "dev"
 Requires-Dist: mypy>=0.800; extra == "dev"
@@ -147,7 +150,7 @@ 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:8081 | xargs kill -9  # dispatcher
 lsof -t -i:3000 | xargs kill -9  # frontend server
 ```
 
@@ -292,7 +295,7 @@ chmod +x ./rapidfireai/start_dev.sh
 # 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:8081 | xargs kill -9 # dispatcher
 lsof -t -i:5002 | xargs kill -9 # mlflow
 lsof -t -i:3000 | xargs kill -9 # frontend
 ```

{rapidfireai-0.10.2rc4 → rapidfireai-0.10.3rc1}/README.md

@@ -88,7 +88,7 @@ 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:8081 | xargs kill -9  # dispatcher
 lsof -t -i:3000 | xargs kill -9  # frontend server
 ```
 
@@ -233,7 +233,7 @@ chmod +x ./rapidfireai/start_dev.sh
 # 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:8081 | xargs kill -9 # dispatcher
 lsof -t -i:5002 | xargs kill -9 # mlflow
 lsof -t -i:3000 | xargs kill -9 # frontend
 ```

{rapidfireai-0.10.2rc4 → rapidfireai-0.10.3rc1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "rapidfireai"
-version = "0.10.2rc4"
+version = "0.10.3rc1"
 authors = [
     {name = "RapidFire AI Inc.", email = "support@rapidfire.ai"},
 ]
@@ -27,6 +27,7 @@ classifiers = [
 ]
 dependencies = [
     "pandas>=2.3.1",
+    "psutil>=5.9.0",
     "torch>=2.8.0",
     "transformers>=4.55.2",
     "peft>=0.17.0",
@@ -45,14 +46,16 @@ dependencies = [
     "loguru>=0.7.3",
     "ipython>=8.37.0",
     "jupyter>=1.1.1",
-    "ipywidgets>=8.1.7", 
+    "ipywidgets>=7.3.4,<9.0.0",  # Support both v7 (Colab) and v8 (Jupyter)
     "uv>=0.8.14",
+    "tensorboard>=2.11.0",
 ]
 
 [project.optional-dependencies]
 dev = [
     "pytest>=6.0",
     "pytest-cov>=2.0",
+    "pytest-mock>=3.10.0",
     "black>=21.0",
     "flake8>=3.8",
     "mypy>=0.800",

{rapidfireai-0.10.2rc4 → rapidfireai-0.10.3rc1}/rapidfireai/backend/controller.py

@@ -9,7 +9,6 @@ from pathlib import Path
 from pprint import pformat
 from typing import Any
 
-import mlflow
 import torch
 from torch.utils.data import Dataset
 
@@ -20,6 +19,7 @@ from rapidfireai.db.rf_db import RfDb
 from rapidfireai.utils.automl_utils import get_flattened_config_leaf, get_runs
 from rapidfireai.utils.constants import (
     MLFLOW_URL,
+    TENSORBOARD_LOG_DIR,
     ControllerTask,
     ExperimentTask,
     RunEndedBy,
@@ -27,11 +27,12 @@ from rapidfireai.utils.constants import (
     RunStatus,
     TaskStatus,
     WorkerTask,
+    get_tracking_backend,
 )
 from rapidfireai.utils.datapaths import DataPath
 from rapidfireai.utils.exceptions import ControllerException, NoGPUsFoundException
 from rapidfireai.utils.logging import RFLogger
-from rapidfireai.utils.mlflow_manager import MLflowManager
+from rapidfireai.utils.metric_logger import create_metric_logger
 from rapidfireai.utils.serialize import encode_payload
 from rapidfireai.utils.shm_manager import SharedMemoryManager
 from rapidfireai.utils.worker_manager import WorkerManager
@@ -79,9 +80,20 @@ class Controller:
             self.num_workers, registry, process_lock
         )
 
-        # create mlflow manager
-        self.mlflow_manager: MLflowManager = MLflowManager(MLFLOW_URL)
-        self.mlflow_manager.get_experiment(self.experiment_name)
+        # create metric logger
+        # Initialize DataPath temporarily to get experiment path for tensorboard logs
+        experiment_path = self.db.get_experiments_path(self.experiment_name)
+        DataPath.initialize(self.experiment_name, experiment_path)
+        tensorboard_log_dir = TENSORBOARD_LOG_DIR or str(DataPath.experiments_path / "tensorboard_logs")
+
+        self.metric_logger = create_metric_logger(
+            backend=get_tracking_backend(),
+            mlflow_tracking_uri=MLFLOW_URL,
+            tensorboard_log_dir=tensorboard_log_dir,
+        )
+        # Get experiment if using MLflow
+        if hasattr(self.metric_logger, 'get_experiment'):
+            self.metric_logger.get_experiment(self.experiment_name)
 
         self.logger.debug("Controller initialized")
 
@@ -148,20 +160,20 @@ class Controller:
                     f"Failed to create required Run DataPath directories: {e}"
                 ) from e
 
-            # create new MlFlow run
+            # create new tracking run
             try:
-                # create new MlFlow run and get the mlflow_run_id
-                mlflow_run_id = self.mlflow_manager.create_run(str(run_id))
+                # create new tracking run and get the mlflow_run_id
+                mlflow_run_id = self.metric_logger.create_run(str(run_id))
 
-                # populate MLFlow with model config info
+                # populate tracking backend with model config info
                 for key, value in flattened_config.items():
-                    self.mlflow_manager.log_param(mlflow_run_id, key, value)
+                    self.metric_logger.log_param(mlflow_run_id, key, value)
                 if warm_started_from:
-                    self.mlflow_manager.log_param(
+                    self.metric_logger.log_param(
                         mlflow_run_id, "warm-start", str(warm_started_from)
                     )
                 if cloned_from:
-                    self.mlflow_manager.log_param(
+                    self.metric_logger.log_param(
                         mlflow_run_id, "parent-run", str(cloned_from)
                     )
                 self.logger.debug(
@@ -172,10 +184,11 @@ class Controller:
                     mlflow_run_id=mlflow_run_id,
                     flattened_config=flattened_config,
                 )
-            except mlflow.exceptions.MlflowException as e:
-                msg = f"Error creating new MLFlow run for run {run_id} - {e}."
+            except Exception as e:
+                # Catch any metric logger exceptions (MLflow, TensorBoard, etc.)
+                msg = f"Error creating new tracking run for run {run_id} - {e}."
                 print(msg)
-                self.mlflow_manager.end_run(mlflow_run_id)
+                self.metric_logger.end_run(mlflow_run_id)
                 self.logger.error(msg, exc_info=True)
 
         total_runs = len(runs)
@@ -249,7 +262,7 @@ class Controller:
 
                 # delete run from MLFlow
                 mlflow_run_id = self.db.get_run(run_id)["mlflow_run_id"]
-                self.mlflow_manager.delete_run(mlflow_run_id)
+                self.metric_logger.delete_run(mlflow_run_id)
                 # mark run as deleted
                 self.db.set_run_details(
                     run_id=run_id,

{rapidfireai-0.10.2rc4 → rapidfireai-0.10.3rc1}/rapidfireai/backend/worker.py

@@ -24,11 +24,11 @@ from rapidfireai.ml.checkpoint_utils import (
     save_model_to_shared_memory,
 )
 from rapidfireai.ml.trainer import create_trainer_instance
-from rapidfireai.utils.constants import MLFLOW_URL, USE_SHARED_MEMORY, RunStatus, SHMObjectType, TaskStatus, WorkerTask
+from rapidfireai.utils.constants import MLFLOW_URL, USE_SHARED_MEMORY, TENSORBOARD_LOG_DIR, RunStatus, SHMObjectType, TaskStatus, WorkerTask, get_tracking_backend
 from rapidfireai.utils.datapaths import DataPath
 from rapidfireai.utils.exceptions import WorkerException
 from rapidfireai.utils.logging import RFLogger, TrainingLogger
-from rapidfireai.utils.mlflow_manager import MLflowManager
+from rapidfireai.utils.metric_logger import create_metric_logger
 from rapidfireai.utils.serialize import decode_db_payload
 from rapidfireai.utils.shm_manager import SharedMemoryManager
 from rapidfireai.utils.trainer_config import TrainerConfig
@@ -71,13 +71,20 @@ class Worker:
         # get experiment name
         self.experiment_name: str = self.db.get_running_experiment()["experiment_name"]
 
-        # create mlflow manager
-        self.mlflow_manager: MLflowManager = MLflowManager(MLFLOW_URL)
-        self.mlflow_manager.get_experiment(self.experiment_name)
-
         # initialize data paths
         DataPath.initialize(self.experiment_name, self.db.get_experiments_path(self.experiment_name))
 
+        # create metric logger
+        tensorboard_log_dir = TENSORBOARD_LOG_DIR or str(DataPath.experiments_path / "tensorboard_logs")
+        self.metric_logger = create_metric_logger(
+            backend=get_tracking_backend(),
+            mlflow_tracking_uri=MLFLOW_URL,
+            tensorboard_log_dir=tensorboard_log_dir,
+        )
+        # Get experiment if using MLflow
+        if hasattr(self.metric_logger, 'get_experiment'):
+            self.metric_logger.get_experiment(self.experiment_name)
+
         # load datasets
         self.train_dataset, self.eval_dataset, self.num_chunks = self.load_datasets()
         self.len_train_dataset = len(self.train_dataset)
@@ -149,7 +156,7 @@ class Worker:
         stderr_buffer = StringIO()
         with redirect_stdout(stdout_buffer), redirect_stderr(stderr_buffer):
             trainer_instance, base_model_name = create_trainer_instance(
-                trainer_config, self.shm_manager, USE_SHARED_MEMORY, self.mlflow_manager, chunk_id
+                trainer_config, self.shm_manager, USE_SHARED_MEMORY, self.metric_logger, chunk_id
             )
 
         # if first time, save checkpoint to disk

{rapidfireai-0.10.2rc4 → rapidfireai-0.10.3rc1}/rapidfireai/cli.py

@@ -398,8 +398,35 @@ def main():
         action="version",
         version=f"RapidFire AI {__version__}"
     )
-    
+
+    parser.add_argument(
+        "--tracking-backend",
+        choices=["mlflow", "tensorboard", "both"],
+        default=os.getenv("RF_TRACKING_BACKEND", "mlflow"),
+        help="Tracking backend to use for metrics (default: mlflow)"
+    )
+
+    parser.add_argument(
+        "--tensorboard-log-dir",
+        default=os.getenv("RF_TENSORBOARD_LOG_DIR", None),
+        help="Directory for TensorBoard logs (default: {experiment_path}/tensorboard_logs)"
+    )
+
+    parser.add_argument(
+        "--colab",
+        action="store_true",
+        help="Run in Colab mode (skips frontend, conditionally starts MLflow based on tracking backend)"
+    )
+
     args = parser.parse_args()
+
+    # Set environment variables from CLI args
+    if args.tracking_backend:
+        os.environ["RF_TRACKING_BACKEND"] = args.tracking_backend
+    if args.tensorboard_log_dir:
+        os.environ["RF_TENSORBOARD_LOG_DIR"] = args.tensorboard_log_dir
+    if args.colab:
+        os.environ["RF_COLAB_MODE"] = "true"
     
     # Handle doctor command separately
     if args.command == "doctor":

{rapidfireai-0.10.2rc4 → rapidfireai-0.10.3rc1}/rapidfireai/db/rf_db.py

@@ -113,7 +113,7 @@ class RfDb:
     def create_experiment(
         self,
         experiment_name: str,
-        mlflow_experiment_id: str,
+        mlflow_experiment_id: str | None,
         config_options: dict[str, Any],
     ) -> int:
         """Create a new experiment"""

{rapidfireai-0.10.2rc4 → rapidfireai-0.10.3rc1}/rapidfireai/db/tables.sql

@@ -2,7 +2,7 @@
 CREATE TABLE IF NOT EXISTS experiments (
     experiment_id INTEGER PRIMARY KEY AUTOINCREMENT,
     experiment_name TEXT NOT NULL,
-    mlflow_experiment_id TEXT NOT NULL,
+    mlflow_experiment_id TEXT,
     config_options TEXT NOT NULL,
     status TEXT NOT NULL,
     current_task TEXT NOT NULL,

{rapidfireai-0.10.2rc4 → rapidfireai-0.10.3rc1}/rapidfireai/dispatcher/dispatcher.py

@@ -30,7 +30,9 @@ class Dispatcher:
         self.app: Flask = Flask(__name__)
 
         # Enable CORS for all routes
-        _ = CORS(self.app, resources={r"/*": {"origins": CORS_ALLOWED_ORIGINS}})
+        # 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()

{rapidfireai-0.10.2rc4 → rapidfireai-0.10.3rc1}/rapidfireai/dispatcher/gunicorn.conf.py

@@ -5,7 +5,7 @@ from rapidfireai.utils.constants import DispatcherConfig
 
 # Other Gunicorn settings...
 bind = f"{DispatcherConfig.HOST}:{DispatcherConfig.PORT}"
-workers = 2
+workers = 1  # Single worker for Colab/single-user environments to save memory
 
 wsgi_app = "rapidfireai.dispatcher.dispatcher:serve_forever()"
 

{rapidfireai-0.10.2rc4 → rapidfireai-0.10.3rc1}/rapidfireai/experiment.py

@@ -17,9 +17,11 @@ from rapidfireai.utils.constants import MLFLOW_URL
 from rapidfireai.utils.exceptions import ExperimentException
 from rapidfireai.utils.experiment_utils import ExperimentUtils
 from rapidfireai.utils.logging import RFLogger
-from rapidfireai.utils.mlflow_manager import MLflowManager
 from rapidfireai.version import __version__
 
+# Note: MLflowManager is imported lazily in get_results() to avoid
+# connection attempts when using tensorboard-only mode
+
 
 class Experiment:
     """Class to manage the entire experiment lifecycle."""
@@ -39,6 +41,7 @@ class Experiment:
         self.experiment_id: int | None = None
         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
         try:
@@ -88,13 +91,68 @@ class Experiment:
         seed: int = 42,
     ) -> None:
         """Run the fit"""
+
+        # 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
+
+        # Detect if running in Google Colab
         try:
-            controller = Controller(self.experiment_id, self.experiment_name)
-            controller.run_fit(param_config, create_model_fn, train_dataset, eval_dataset, num_chunks, seed)
-        except Exception as e:
-            if hasattr(self, "logger"):
-                self.logger.opt(exception=True).error(f"Error running fit: {e}")
-            raise ExperimentException(f"Error running fit: {e}, traceback: {traceback.format_exc()}") from e
+            import google.colab
+            in_colab = True
+        except ImportError:
+            in_colab = False
+
+        if in_colab:
+            # Run Controller in background thread to keep kernel responsive
+            import threading
+            import sys
+            from io import StringIO
+            from IPython.display import display, HTML
+
+            def _run_controller_background():
+                """Run controller in background thread with output suppression"""
+                # Suppress stdout to avoid print statements appearing in wrong cells
+                old_stdout = sys.stdout
+                sys.stdout = StringIO()
+
+                try:
+                    controller = Controller(self.experiment_id, self.experiment_name)
+                    controller.run_fit(param_config, create_model_fn, train_dataset, eval_dataset, num_chunks, seed)
+                except Exception as e:
+                    # 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}")
+                    display(HTML(f'<p style="color: red; font-weight: bold;">❌ Error in background training: {e}</p>'))
+                finally:
+                    # Restore stdout
+                    sys.stdout = old_stdout
+                    # Display completion message
+                    display(HTML('<p style="color: blue; font-weight: bold;">🎉 Training completed! Check InteractiveController for final results.</p>'))
+                    self._training_thread = None
+
+            self._training_thread = threading.Thread(target=_run_controller_background, daemon=True)
+            self._training_thread.start()
+
+            # Use IPython display for reliable output in Colab
+            display(HTML(
+                '<div style="padding: 10px; background-color: #d4edda; border: 1px solid #28a745; border-radius: 5px; color: #155724;">'
+                '<b>✓ Training started in background</b><br>'
+                'Use InteractiveController to monitor progress. The notebook kernel will remain responsive while training runs.<br>'
+                '<small>Tip: Interact with InteractiveController periodically to keep Colab active.</small>'
+                '</div>'
+            ))
+        else:
+            # Original blocking behavior for non-Colab environments
+            try:
+                controller = Controller(self.experiment_id, self.experiment_name)
+                controller.run_fit(param_config, create_model_fn, train_dataset, eval_dataset, num_chunks, seed)
+            except Exception as e:
+                if hasattr(self, "logger"):
+                    self.logger.opt(exception=True).error(f"Error running fit: {e}")
+                raise ExperimentException(f"Error running fit: {e}, traceback: {traceback.format_exc()}") from e
 
     def get_results(self) -> pd.DataFrame:
         """
@@ -102,6 +160,16 @@ class Experiment:
         """
         try:
             runs_info_df = self.experiment_utils.get_runs_info()
+
+            # Check if there are any mlflow_run_ids before importing MLflow
+            has_mlflow_runs = runs_info_df.get("mlflow_run_id") is not None and runs_info_df["mlflow_run_id"].notna().any()
+
+            if not has_mlflow_runs:
+                # No MLflow runs to fetch, return empty DataFrame
+                return pd.DataFrame(columns=["run_id", "step"])
+
+            # Lazy import - only import when we actually have MLflow runs to fetch
+            from rapidfireai.utils.mlflow_manager import MLflowManager
             mlflow_manager = MLflowManager(MLFLOW_URL)
 
             metrics_data = []

{rapidfireai-0.10.2rc4 → rapidfireai-0.10.3rc1}/rapidfireai/frontend/build/asset-manifest.json

@@ -1,7 +1,7 @@
 {
   "files": {
     "main.css": "/static-files/static/css/main.702595df.css",
-    "main.js": "/static-files/static/js/main.3ff1e37d.js",
+    "main.js": "/static-files/static/js/main.e7d3b759.js",
     "ml-model-trace-renderer.js": "/static-files/lib/notebook-trace-renderer/js/ml-model-trace-renderer.5490ebc325fe0f300ad9.js",
     "static/js/6019.9025341e.chunk.js": "/static-files/static/js/6019.9025341e.chunk.js",
     "static/js/6336.8153bc1c.chunk.js": "/static-files/static/js/6336.8153bc1c.chunk.js",
@@ -120,7 +120,7 @@
     "static/media/chart-line.svg": "/static-files/static/media/chart-line.0adaa2036bb4eb5956db6d0c7e925a3d.svg",
     "lib/notebook-trace-renderer/index.html": "/static-files/lib/notebook-trace-renderer/index.html",
     "main.702595df.css.map": "/static-files/static/css/main.702595df.css.map",
-    "main.3ff1e37d.js.map": "/static-files/static/js/main.3ff1e37d.js.map",
+    "main.e7d3b759.js.map": "/static-files/static/js/main.e7d3b759.js.map",
     "ml-model-trace-renderer.js.map": "/static-files/lib/notebook-trace-renderer/js/ml-model-trace-renderer.5490ebc325fe0f300ad9.js.map",
     "6336.8153bc1c.chunk.js.map": "/static-files/static/js/6336.8153bc1c.chunk.js.map",
     "9478.cbf55ef3.chunk.js.map": "/static-files/static/js/9478.cbf55ef3.chunk.js.map",
@@ -216,6 +216,6 @@
   },
   "entrypoints": [
     "static/css/main.702595df.css",
-    "static/js/main.3ff1e37d.js"
+    "static/js/main.e7d3b759.js"
   ]
 }

{rapidfireai-0.10.2rc4 → rapidfireai-0.10.3rc1}/rapidfireai/frontend/build/index.html

@@ -1 +1 @@
-<!doctype html><html lang="en"><head><meta charset="utf-8"/><meta name="viewport" content="width=device-width,initial-scale=1,shrink-to-fit=no"/><link rel="shortcut icon" href="./static-files/favicon.ico"/><meta name="theme-color" content="#000000"/><link rel="manifest" href="./static-files/manifest.json" crossorigin="use-credentials"/><title>RapidFire AI</title><script defer="defer" src="static-files/static/js/main.3ff1e37d.js"></script><link href="static-files/static/css/main.702595df.css" rel="stylesheet"></head><body><noscript>You need to enable JavaScript to run this app.</noscript><div id="root" class="mlflow-ui-container"></div><div id="modal" class="mlflow-ui-container"></div></body></html>
+<!doctype html><html lang="en"><head><meta charset="utf-8"/><meta name="viewport" content="width=device-width,initial-scale=1,shrink-to-fit=no"/><link rel="shortcut icon" href="./static-files/favicon.ico"/><meta name="theme-color" content="#000000"/><link rel="manifest" href="./static-files/manifest.json" crossorigin="use-credentials"/><title>RapidFire AI</title><script defer="defer" src="static-files/static/js/main.e7d3b759.js"></script><link href="static-files/static/css/main.702595df.css" rel="stylesheet"></head><body><noscript>You need to enable JavaScript to run this app.</noscript><div id="root" class="mlflow-ui-container"></div><div id="modal" class="mlflow-ui-container"></div></body></html>