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

rapidfireai/backend/scheduler.py

@@ -0,0 +1,137 @@
+"""This module contains the Scheduler class which is responsible for scheduling runs on workers to train on a chunk."""
+
+
+class Scheduler:
+    """This class is responsible for scheduling runs on to workers to train on a chunk"""
+
+    def __init__(self, run_ids: list[int], num_workers: int, num_chunks: int) -> None:
+        # run_ids are 1 indexed
+        # worker_ids are 0 indexed
+        # chunk_ids are 0 indexed
+
+        self.n_runs: int = len(run_ids)
+        self.n_workers: int = num_workers
+        self.n_chunks: int = num_chunks
+        self.run_ids: list[int] = run_ids
+
+        # create data structures
+        self.worker_running_current_run: dict[int, int] = dict.fromkeys(range(self.n_workers), -1)
+        self.run_visited_num_chunks: dict[int, int] = dict.fromkeys(self.run_ids, 0)
+        self.run_start_chunk_id: dict[int, int] = dict.fromkeys(self.run_ids, 0)
+
+        # add runs to scheduler
+        for run_id in run_ids:
+            self.add_run(run_id, 0)
+
+    def reset_run(self, run_id: int) -> None:
+        """Reset the scheduler for a specific run (used at epoch boundaries)"""
+        if run_id in self.run_ids:
+            # Reset progress for this run
+            self.run_visited_num_chunks[run_id] = 0
+
+            # If this run is currently assigned to a worker, free the worker
+            for worker_id in range(self.n_workers):
+                if self.worker_running_current_run[worker_id] == run_id:
+                    self.worker_running_current_run[worker_id] = -1
+
+    def add_run(self, run_id: int, run_visited_num_chunks: int, run_start_chunk_id: int = 0) -> None:
+        """Add a new run to the scheduler."""
+        if run_id not in self.run_ids:
+            self.run_ids.append(run_id)
+            self.n_runs = len(self.run_ids)
+
+        self.run_visited_num_chunks[run_id] = run_visited_num_chunks
+        self.run_start_chunk_id[run_id] = run_start_chunk_id
+
+    def set_completed_task(self, worker_id: int) -> None:
+        """Set a task as completed."""
+        run_id = self.worker_running_current_run[worker_id]
+
+        if run_id != -1:
+            self.worker_running_current_run[worker_id] = -1
+            self.run_visited_num_chunks[run_id] += 1
+
+    def remove_run(self, run_id: int) -> int:
+        """Remove a run from the scheduler and return its progress."""
+        if run_id not in self.run_ids:
+            return 0
+
+        # Get the progress before removing
+        progress = self.run_visited_num_chunks.get(run_id, 0)
+
+        # Clean up worker assignment
+        for worker_id in range(self.n_workers):
+            if self.worker_running_current_run[worker_id] == run_id:
+                self.worker_running_current_run[worker_id] = -1
+
+        # Remove from all data structures
+        self.run_visited_num_chunks.pop(run_id, None)
+        self.run_start_chunk_id.pop(run_id, None)
+
+        if run_id in self.run_ids:
+            self.run_ids.remove(run_id)
+            self.n_runs = len(self.run_ids)
+
+        return progress
+
+    def schedule(self) -> dict[str, int | bool | None] | None:
+        """
+        Schedule a single task based on constraints and preferences.
+        Returns {run_id: <>, worker_id: <>, chunk_id: <>, is_last_chunk: <>} if a schedule is possible.
+        Returns {run_id: None, worker_id: None, chunk_id: None, is_last_chunk: None} if all runs have seen all chunks.
+        Returns {run_id: -1, worker_id: -1, chunk_id: -1, is_last_chunk: None} if all workers are busy or no runs are available.
+        """
+        # First check if all workers are busy (most common condition)
+        available_workers = [
+            worker_id for worker_id in range(self.n_workers) if self.worker_running_current_run[worker_id] == -1
+        ]
+        if not available_workers:
+            return {"run_id": -1, "worker_id": -1, "chunk_id": -1, "is_last_chunk": None}
+
+        # Next check if all runs have seen all chunks (termination condition)
+        if all(self.run_visited_num_chunks[run_id] >= self.n_chunks for run_id in self.run_ids):
+            return {"run_id": None, "worker_id": None, "chunk_id": None, "is_last_chunk": None}
+
+        # Get busy runs and available runs
+        busy_runs = set(run_id for run_id in self.worker_running_current_run.values() if run_id != -1)
+        available_runs = [
+            run_id
+            for run_id in self.run_ids
+            if self.run_visited_num_chunks[run_id] < self.n_chunks and run_id not in busy_runs
+        ]
+
+        # If no available runs, return busy state
+        if not available_runs:
+            return {"run_id": -1, "worker_id": -1, "chunk_id": -1, "is_last_chunk": None}
+
+        # Find the run with least progress, then lowest run_id for tie-breaking
+        # NOTE: any newly inserted clones will take priority
+        # NOTE: prioritize by run_id if run_visited_num_chunks is the same
+
+        run_id = min(available_runs, key=lambda run_id: (self.run_visited_num_chunks[run_id], run_id))
+        worker_id = available_workers[0]  # Pick first available worker
+        chunk_id = (
+            self.run_visited_num_chunks[run_id] + self.run_start_chunk_id[run_id]
+        ) % self.n_chunks  # Next chunk in sequence starting from run_start_chunk_id (chunk_id is 0-indexed)
+        is_last_chunk = chunk_id == self.n_chunks - 1
+
+        # Update internal state immediately
+        self.worker_running_current_run[worker_id] = run_id
+
+        return {"run_id": run_id, "worker_id": worker_id, "chunk_id": chunk_id, "is_last_chunk": is_last_chunk}
+
+    def get_status(self) -> dict:
+        """Get current scheduler status for debugging."""
+        completed_runs = [run_id for run_id in self.run_ids if self.run_visited_num_chunks[run_id] == self.n_chunks]
+
+        return {
+            "active_runs": len([r for r in self.run_ids if self.run_visited_num_chunks[r] < self.n_chunks]),
+            "busy_workers": len([w for w in range(self.n_workers) if self.worker_running_current_run[w] != -1]),
+            "completed_runs": len(completed_runs),
+            "worker_assignments": {
+                w: self.worker_running_current_run[w]
+                for w in range(self.n_workers)
+                if self.worker_running_current_run[w] != -1
+            },
+            "run_progress": {r: f"{self.run_visited_num_chunks[r]}/{self.n_chunks}" for r in self.run_ids},
+        }

rapidfireai/backend/worker.py

@@ -0,0 +1,272 @@
+"""This module contains the Worker class which is responsible for handling the worker operations."""
+
+import gc
+import os
+import time
+import traceback
+from contextlib import redirect_stderr, redirect_stdout
+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, Callable
+
+import torch
+
+from rapidfireai.backend.chunks import DatasetChunks
+from rapidfireai.db.rf_db import RfDb
+from rapidfireai.ml.checkpoint_utils import (
+    save_checkpoint_to_disk,
+    save_checkpoint_to_shared_memory,
+    save_model_to_shared_memory,
+)
+from rapidfireai.ml.trainer import create_trainer_instance
+from rapidfireai.utils.constants import MLFLOW_URL, RunStatus, TaskStatus, WorkerTask
+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.serialize import decode_db_payload
+from rapidfireai.utils.shm_manager import SharedMemoryManager
+from rapidfireai.utils.trainer_config import TrainerConfig
+
+
+class Worker:
+    """Worker class that handles training and validation of runs"""
+
+    def __init__(self, worker_id: int, model_registry: DictProxy, process_lock: Lock, shutdown_event: EventType):
+        """Initialize the worker"""
+        self.process: Process
+        self.worker_id: int = worker_id
+        self.shutdown_event: EventType = shutdown_event
+
+        # Shared memory attributes (set by WorkerManager)
+        self.model_registry: DictProxy[int, Any] = model_registry
+        self.process_lock: Lock = process_lock
+
+        # Shared memory manager will be created using global objects
+        self.shm_manager = SharedMemoryManager(
+            name=f"worker-{worker_id}-shm", registry=model_registry, multiprocess_lock=process_lock
+        )
+
+        # 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"]
+
+        # 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))
+
+        # load datasets
+        train_dataset, self.eval_dataset, self.num_chunks = self.load_datasets()
+        self.len_train_dataset = len(train_dataset)
+        self.train_dataset_chunks = DatasetChunks(train_dataset, self.num_chunks)
+
+    def load_datasets(self) -> tuple[torch.utils.data.Dataset | None, torch.utils.data.Dataset | None, int]:
+        """Load the train and eval datasets"""
+        try:
+            with open(DataPath.dataset_path(), "rb") as f:
+                datasets = decode_db_payload(f.read())
+            self.logger.debug("Loaded datasets")
+            return datasets["train"], datasets["eval"], datasets["num_chunks"]
+        except Exception as e:
+            raise WorkerException(f"Error loading datasets: {e}") from e
+
+    def run_fit(
+        self,
+        run_id: int,
+        chunk_id: int,
+        create_model_fn: Callable,
+    ) -> None:
+        """Run fit"""
+        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"]
+        mlflow_run_id = run_details["mlflow_run_id"]
+
+        # set seed
+        # torch.manual_seed(run_details["seed"])
+        # np.random.seed(run_details["seed"])
+        # random.seed(run_details["seed"])
+
+        # fetch train dataset chunk
+        train_dataset_chunk = self.train_dataset_chunks.get_chunk(chunk_id)
+        # create worker config
+        trainer_config = TrainerConfig(
+            worker_id=self.worker_id,
+            run_id=run_id,
+            mlflow_run_id=mlflow_run_id,
+            config_leaf=config_leaf,
+            total_steps=run_details["total_steps"],
+            completed_steps=run_details["completed_steps"],
+            create_model_fn=create_model_fn,
+            train_dataset=train_dataset_chunk,
+            eval_dataset=self.eval_dataset,
+            warm_started_from=run_details["warm_started_from"],
+            num_epochs_completed=run_details["num_epochs_completed"],
+        )
+        completed_steps = self.db.get_completed_steps(run_id)
+
+        use_shared_memory = True
+        parent_run_details = None
+        if trainer_config.warm_started_from is not None:
+            parent_run_details = self.db.get_run(trainer_config.warm_started_from)
+            if trainer_config.config_leaf.get("trainer_type") == "GRPO":
+                config_leaf["reward_funcs"] = parent_run_details["config_leaf"].get("reward_funcs")
+                self.db.set_run_details(run_id, config_leaf=config_leaf)
+
+        stdout_buffer = StringIO()
+        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
+            )
+
+        # if first time, save checkpoint to disk
+        if completed_steps == 0 and not use_shared_memory:
+            save_checkpoint_to_disk(trainer_instance, trainer_config, first=True)
+
+        # write logs to user logger
+        if stdout_buffer.getvalue():
+            self.training_logger.info(stdout_buffer.getvalue())
+        if stderr_buffer.getvalue():
+            self.training_logger.error(stderr_buffer.getvalue())
+
+        self.logger.debug(f"Beginning training for run {run_id} on chunk {chunk_id}")
+
+        # update base model name in db for run
+        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)
+
+        # Train the model
+        stdout_buffer = StringIO()
+        stderr_buffer = StringIO()
+        with redirect_stdout(stdout_buffer), redirect_stderr(stderr_buffer):
+            trainer_instance.train()
+
+        # write logs to user logger
+        if stdout_buffer.getvalue():
+            self.training_logger.info(stdout_buffer.getvalue())
+        if stderr_buffer.getvalue():
+            self.training_logger.error(stderr_buffer.getvalue())
+
+        # update completed steps
+        new_completed_steps = completed_steps + trainer_instance.state.global_step
+        self.db.set_completed_steps(run_id, new_completed_steps)
+
+        save_strategy = trainer_config.config_leaf.get("training_args", {}).get("save_strategy", "epoch")
+
+        # Save checkpoints to shared memory
+        if use_shared_memory:
+            save_checkpoint_to_shared_memory(trainer_instance, trainer_config, self.shm_manager)
+            if not trainer_config.config_leaf.get("peft_params"):
+                save_model_to_shared_memory(
+                    trainer_instance.model,
+                    trainer_instance.tokenizer,
+                    trainer_config,
+                    self.shm_manager,
+                    "full_model",
+                    trainer_config.run_id,
+                )
+            self.logger.debug(f"Saved checkpoint to shared memory for run {run_id} on chunk {chunk_id}")
+            if save_strategy == "chunk" or (save_strategy == "epoch" and chunk_id == self.num_chunks - 1):
+                save_checkpoint_to_disk(trainer_instance, trainer_config, completed_steps=new_completed_steps)
+                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)
+            self.logger.debug(f"Saved checkpoint to disk for run {run_id} on chunk {chunk_id}")
+
+        if chunk_id == self.num_chunks - 1 and new_completed_steps >= trainer_config.total_steps:
+            save_checkpoint_to_disk(trainer_instance, trainer_config, last=True)
+            self.logger.debug(f"Saved final checkpoint for run {run_id} on chunk {chunk_id}")
+
+        # clean up all references to shared memory objects
+        if hasattr(trainer_instance, "model"):
+            del trainer_instance.model
+        if hasattr(trainer_instance, "ref_model"):
+            del trainer_instance.ref_model
+        if hasattr(trainer_instance, "optimizer"):
+            del trainer_instance.optimizer
+        if hasattr(trainer_instance, "lr_scheduler"):
+            del trainer_instance.lr_scheduler
+        del trainer_instance
+
+        # run garbage collection
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        self.logger.debug(f"Completed training for run {run_id} on chunk {chunk_id}")
+
+    def serve_forever(self) -> None:
+        """This runs in the worker process"""
+
+        prev_task_id: int | None = None
+        while not (self.shutdown_event and self.shutdown_event.is_set()):
+            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"]
+                chunk_id = scheduled_task["chunk_id"]
+                create_model_fn = scheduled_task["config_options"]["create_model_fn"]
+                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)
+
+                    # run train and validation function
+                    try:
+                        self.run_fit(run_id, chunk_id, create_model_fn)
+                        self.db.set_worker_task_status(self.worker_id, TaskStatus.COMPLETED)
+                    except Exception as e:
+                        self.logger.opt(exception=True).error(
+                            f"Error while running run_fit for run {run_id} and chunk {chunk_id}: {e}"
+                        )
+                        self.db.set_run_details(run_id, status=RunStatus.FAILED, error=str(e) + traceback.format_exc())
+                        self.db.set_worker_task_status(self.worker_id, TaskStatus.FAILED)
+                else:
+                    raise WorkerException(f"Invalid task type: {task_type}")
+            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())
+                break
+
+        self.shutdown()
+
+    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()
+
+        # 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()