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

rapidfireai/backend/controller.py

@@ -0,0 +1,637 @@
+"""This module contains the Controller class which is responsible for orchestrating the RapidFire lifecycle."""
+
+import math
+import random
+import time
+from logging import Logger
+from pathlib import Path
+from pprint import pformat
+from typing import Any, Callable
+
+import mlflow
+import torch
+from torch.utils.data import Dataset
+
+from rapidfireai.automl import AutoMLAlgorithm
+from rapidfireai.backend.scheduler import Scheduler
+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,
+    ControllerTask,
+    ExperimentTask,
+    RunEndedBy,
+    RunSource,
+    RunStatus,
+    TaskStatus,
+    WorkerTask,
+)
+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.serialize import encode_payload
+from rapidfireai.utils.shm_manager import SharedMemoryManager
+from rapidfireai.utils.worker_manager import WorkerManager
+
+
+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
+
+        try:
+            mp.set_start_method("spawn", force=True)
+        except RuntimeError:
+            # Start method already set
+            pass
+
+        self.experiment_id: int = experiment_id
+        self.experiment_name: str = experiment_name
+
+        # 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")
+
+        # 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.")
+
+        # 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()
+
+        # create worker manager
+        self.worker_manager: WorkerManager = WorkerManager(self.num_workers, registry, process_lock)
+
+        # create mlflow manager
+        self.mlflow_manager: MLflowManager = MLflowManager(MLFLOW_URL)
+        self.mlflow_manager.get_experiment(self.experiment_name)
+
+        self.logger.debug("Controller initialized")
+
+    def _create_models(
+        self,
+        param_config: AutoMLAlgorithm | dict[str, Any],
+        source: RunSource,
+        seed: int,
+        len_train_dataset: int,
+        num_chunks: int,
+        warm_start_info: dict[str, Any] | None = None,
+    ) -> list[int]:
+        """Create the models."""
+
+        # get config_leaf from param_config for each run
+        config_leafs = get_runs(param_config, seed)
+
+        # create runs
+        runs = {}
+        for config_leaf in config_leafs:
+            flattened_config = get_flattened_config_leaf(config_leaf)
+            # print("flattened_config: ",flattened_config)
+            total_steps = self._get_total_step(config_leaf, len_train_dataset, num_chunks)
+
+            run_id = self.db.create_run(
+                config_leaf=config_leaf,
+                status=RunStatus.NEW,
+                completed_steps=0,
+                total_steps=total_steps,
+                error="",
+                source=source,
+                ended_by=None,
+                start_chunk_id=warm_start_info["start_chunk_id"] if warm_start_info else 0,
+                warm_started_from=warm_start_info["parent_run_id"] if warm_start_info else None,
+            )
+            runs[run_id] = flattened_config
+
+            # create directories for each run
+            try:
+                base_run_path = DataPath.base_run_path(run_id)
+                work_dir_path = DataPath.work_dir_path(base_run_path)
+                initial_checkpoint_path = DataPath.initial_checkpoint_path(base_run_path)
+                final_checkpoint_path = DataPath.final_checkpoint_path(base_run_path)
+                intermediate_checkpoint_path = DataPath.intermediate_checkpoint_path(base_run_path)
+
+                Path.mkdir(work_dir_path, parents=True, exist_ok=True)
+                Path.mkdir(initial_checkpoint_path, parents=True, exist_ok=True)
+                Path.mkdir(final_checkpoint_path, parents=True, exist_ok=True)
+                Path.mkdir(intermediate_checkpoint_path, parents=True, exist_ok=True)
+            except (PermissionError, OSError) as e:
+                raise ControllerException(f"Failed to create required Run DataPath directories: {e}") from e
+
+            # create new MlFlow run
+            try:
+                # create new MlFlow run and get the mlflow_run_id
+                mlflow_run_id = self.mlflow_manager.create_run(str(run_id))
+
+                # populate MLFlow with model config info
+                for key, value in flattened_config.items():
+                    self.mlflow_manager.log_param(mlflow_run_id, key, value)
+                if warm_start_info:
+                    self.mlflow_manager.log_param(mlflow_run_id, "warm-start", str(warm_start_info))
+                self.logger.debug(f"Populated MLFlow with model config info for run {run_id}.")
+                self.db.set_run_details(
+                    run_id=run_id,
+                    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}."
+                print(msg)
+                self.mlflow_manager.end_run(mlflow_run_id)
+                self.logger.error(msg, exc_info=True)
+
+        total_runs = len(runs)
+        self.logger.info(f"Created {total_runs} runs - \n{pformat(runs, indent=4, width=120)}")
+        self.logger.debug(f"Got {total_runs} runs for {source.value}.")
+
+        # set experiment task to run_fit
+        self.db.set_experiment_current_task(ExperimentTask.RUN_FIT)
+        self.logger.debug("Completed creating models.")
+
+        return list(runs.keys())
+
+    def _clear_run_from_shm(self, run_id: int) -> None:
+        """Clear the run from shared memory."""
+
+        # check if there are any other runs with the same base model
+        base_model_name = self.db.get_run(run_id)["config_leaf"]["model_name"]
+        relevant_runs = self.db.get_runs_by_status([RunStatus.ONGOING, RunStatus.NEW, RunStatus.STOPPED])
+
+        # get shared object types to delete - if no other runs are using it
+        delete_shared_objects = True
+        for r_run_id, r_run_details in relevant_runs.items():
+            if r_run_details["config_leaf"]["model_name"] == base_model_name and r_run_id != run_id:
+                delete_shared_objects = False
+                break
+
+        # delete model object from shared memory
+        self.shm_manager.delete_model_object(run_id, base_model_name if delete_shared_objects else None)
+
+    def _process_interactive_control(
+        self,
+        run_states: dict[str, Any],
+        clone_modify_tasks: list[dict[str, Any]],
+        len_train_dataset: int,
+        seed: int,
+        num_chunks: int,
+    ) -> None:
+        """Process interactive control tasks."""
+
+        # process non-clone_modify tasks
+        for run_id, run_state in run_states.items():
+            if not run_state["task_id"]:
+                continue
+
+            if run_state["status"] == RunStatus.STOPPED:
+                # process stopped tasks
+                # mark run as stopped
+                self.db.set_run_details(
+                    run_id=run_id,
+                    status=RunStatus.STOPPED,
+                    ended_by=RunEndedBy.INTERACTIVE_CONTROL,
+                )
+                self.db.set_ic_ops_task_status(run_state["task_id"], TaskStatus.COMPLETED)
+                self.ic_logger.info(f"Stopping run {run_id} by Interactive Control")
+            elif run_state["status"] == RunStatus.DELETED:
+                # process deleted tasks
+                # clear run from shm
+                self._clear_run_from_shm(run_id)
+                # delete run from MLFlow
+                mlflow_run_id = self.db.get_run(run_id)["mlflow_run_id"]
+                self.mlflow_manager.delete_run(mlflow_run_id)
+                # mark run as deleted
+                self.db.set_run_details(
+                    run_id=run_id,
+                    status=RunStatus.DELETED,
+                    ended_by=RunEndedBy.INTERACTIVE_CONTROL,
+                )
+                self.db.set_ic_ops_task_status(run_state["task_id"], TaskStatus.COMPLETED)
+                self.ic_logger.info(f"Deleting run {run_id} by Interactive Control")
+            elif run_state["status"] == RunStatus.ONGOING:
+                # process ongoing tasks
+                self.db.set_run_details(
+                    run_id=run_id,
+                    status=RunStatus.ONGOING,
+                    ended_by="",
+                )
+                self.db.set_ic_ops_task_status(run_state["task_id"], TaskStatus.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)
+            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"],
+            )
+
+            # add additional_kwargs to config_leaf if it exists in the parent run
+            parent_run_details = self.db.get_run(parent_run_id)
+            if "additional_kwargs" in parent_run_details["config_leaf"]:
+                config_leaf["additional_kwargs"] = parent_run_details["config_leaf"]["additional_kwargs"]
+
+            # create model for the new run
+            try:
+                if ic_op == ControllerTask.IC_CLONE_MODIFY:
+                    run_ids = self._create_models(
+                        config_leaf,
+                        RunSource.INTERACTIVE_CONTROL,
+                        seed,
+                        len_train_dataset,
+                        num_chunks=num_chunks,
+                    )
+                elif ic_op == ControllerTask.IC_CLONE_MODIFY_WARM:
+                    warm_start_info = {
+                        "parent_run_id": parent_run_id,
+                        "start_chunk_id": parent_run_details["num_chunks_visited_curr_epoch"],
+                    }
+                    run_ids = self._create_models(
+                        config_leaf,
+                        RunSource.INTERACTIVE_CONTROL,
+                        seed,
+                        len_train_dataset,
+                        num_chunks,
+                        warm_start_info,
+                    )
+                else:
+                    raise ValueError(f"Unsupported IC operation {ic_op}")
+
+                # mark task as completed
+                self.db.set_ic_ops_task_status(task["task_id"], TaskStatus.COMPLETED)
+                self.ic_logger.info(
+                    f"Cloned run {parent_run_id} by Interactive Control with {ic_op.value} 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
+
+    def _process_interm_ic_ops_states(
+        self,
+        currently_scheduled_runs: list[int],
+    ) -> 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()
+
+        # 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"]
+
+            # 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,
+            )
+
+            if is_clone_modify_task:
+                # clone_modify tasks
+                # get latest run state
+                run_status = run_states[run_id]["status"] if run_id in run_states else self.db.get_run(run_id)["status"]
+
+                # 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}.")
+                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}.")
+            else:
+                # Non clone_modify tasks
+                if run_id not in run_states:
+                    run_states[run_id] = {
+                        "task_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,
+                ]:
+                    # 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)
+                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)
+                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)
+
+                    # add new task to run states
+                    if task["ic_op"] == ControllerTask.IC_STOP:
+                        updated_status = RunStatus.STOPPED
+                        info_msg = f"Received STOP task for run {run_id}"
+                    elif task["ic_op"] == ControllerTask.IC_DELETE:
+                        updated_status = RunStatus.DELETED
+                        info_msg = f"Received DELETE task for run {run_id}"
+                    elif task["ic_op"] == ControllerTask.IC_RESUME:
+                        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']}")
+                    run_states[run_id].update(
+                        {
+                            "task_id": task["task_id"],
+                            "task": task["ic_op"],
+                            "status": updated_status if updated_status else current_status,
+                        }
+                    )
+                    self.ic_logger.info(info_msg)
+
+        return run_states, clone_modify_tasks
+
+    def _get_total_step(self, config_leaf: dict[str, Any], len_train_dataset: int, num_chunks: int) -> int:
+        """Get the total number of steps for a run."""
+        num_train_epochs = config_leaf["training_args"].get("num_train_epochs", 1)
+
+        total_steps = 0
+        # max_steps overrides num_train_epochs
+        if config_leaf["training_args"].get("max_steps", None):
+            # ceil to nearest chunk multiple
+            total_steps = config_leaf["training_args"]["max_steps"]
+        elif num_train_epochs:
+            # total_steps = num_epochs to num_steps =
+            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)
+            total_steps = (
+                math.ceil(len_train_dataset / (num_chunks * per_device_train_batch_size * gradient_accumulation_steps))
+                * num_chunks
+                * num_train_epochs
+            )
+        return total_steps
+
+    def run_fit(
+        self,
+        param_config: Any,
+        create_model_fn: Callable,
+        train_dataset: Dataset,
+        eval_dataset: Dataset,
+        num_chunks: int,
+        seed: int = 42,
+    ) -> None:
+        """Run the fit."""
+
+        # set experiment task to create models
+        self.db.set_experiment_current_task(ExperimentTask.CREATE_MODELS)
+        self.logger.debug(f"Set experiment task to {ExperimentTask.CREATE_MODELS.value}.")
+
+        # save train and eval dataset objects to a file for workers to load
+        try:
+            datasets = {
+                "train": train_dataset,
+                "eval": eval_dataset if eval_dataset else None,
+                "num_chunks": num_chunks,
+            }
+            with open(DataPath.dataset_path(), "w", encoding="utf-8") as f:
+                f.write(encode_payload(datasets))
+            self.logger.debug(f"Saved datasets to {DataPath.dataset_path()}")
+        except Exception as e:
+            raise ControllerException(f"Error saving datasets: {e}") from e
+
+        # set seed
+        random.seed(seed)
+        self.logger.info(f"Set seed to {seed}")
+
+        # create models
+        try:
+            len_train_dataset = len(train_dataset)
+            self._create_models(param_config, RunSource.INITIAL, seed, len_train_dataset, num_chunks=num_chunks)
+            self.logger.debug("Created models.")
+        except Exception as e:
+            raise ControllerException(f"Error creating models: {e}") from e
+
+        # set experiment task to create models
+        self.db.set_experiment_current_task(ExperimentTask.RUN_FIT)
+        self.logger.debug(f"Set experiment task to {ExperimentTask.RUN_FIT.value}.")
+
+        # create workers
+        try:
+            self.worker_manager.create_workers()
+            print("Created workers")
+            self.logger.debug(f"Created {self.num_workers} workers.")
+        except Exception as e:
+            raise ControllerException(f"Error creating workers: {e}") from e
+
+        # create scheduler
+        run_ids = list(
+            self.db.get_runs_by_status(
+                [
+                    RunStatus.NEW,
+                ]
+            ).keys()
+        )
+        scheduler = Scheduler(run_ids, self.num_workers, num_chunks)
+
+        # run fit
+        self.logger.info("Starting Training and Validation")
+        try:
+            all_done = False
+            prev_worker_tasks = {}  # Track previous iteration's worker tasks
+
+            while not all_done:
+                # check for errors
+                exp_error = self.db.get_experiment_error()
+                if exp_error:
+                    print(f"Error in experiment: {exp_error}")
+                    self.logger.error(f"Error in experiment: {exp_error}")
+                    break
+
+                # get current state (pre IC ops states)
+                all_worker_tasks = self.db.get_all_worker_tasks()
+                all_run_details = self.db.get_all_runs()
+
+                # Filter and separate fresh completed and failed tasks in a single loop
+                completed_tasks = {}
+                failed_tasks = []
+                for worker_id, worker_task in all_worker_tasks.items():
+                    prev_task = prev_worker_tasks.get(worker_id, {})
+                    current_task_tuple = (worker_task["task_id"], worker_task["status"])
+                    prev_task_tuple = (prev_task.get("task_id"), prev_task.get("status"))
+
+                    # skip if task is the same as previous iteration (no change in status) or run is not active
+                    if current_task_tuple == prev_task_tuple or worker_task["run_id"] not in scheduler.run_ids:
+                        continue
+
+                    if worker_task["status"] == TaskStatus.COMPLETED:
+                        completed_tasks[worker_id] = worker_task
+                    elif worker_task["status"] == TaskStatus.FAILED:
+                        failed_tasks.append(worker_task)
+
+                # Process completed tasks first (before scheduling new ones)
+                for worker_id, worker_task in completed_tasks.items():
+                    run_id = worker_task["run_id"]
+                    chunk_id = worker_task["chunk_id"]
+                    run_details = all_run_details[run_id]
+                    self.logger.debug(f"Completed task: run {run_id}, chunk {chunk_id} on worker {worker_id}")
+                    self.logger.info(
+                        f"Run {run_id} completed steps - {run_details['completed_steps']}/{run_details['total_steps']}"
+                    )
+
+                    # Update scheduler state
+                    scheduler.set_completed_task(worker_id)
+
+                    # Update database state and local state using scheduler's state as source of truth
+                    new_chunks_visited = scheduler.run_visited_num_chunks[run_id]
+                    if new_chunks_visited == num_chunks:
+                        num_epochs_completed = run_details["num_epochs_completed"] + 1
+                    else:
+                        num_epochs_completed = run_details["num_epochs_completed"]
+                    self.db.set_run_details(
+                        run_id=run_id,
+                        num_chunks_visited_curr_epoch=new_chunks_visited,
+                        num_epochs_completed=num_epochs_completed,
+                    )
+
+                    # Update progress
+                    progress_percentage = (
+                        (run_details["completed_steps"] / run_details["total_steps"] * 100)
+                        if run_details["total_steps"] > 0
+                        else 0
+                    )
+                    self.db.set_controller_progress(run_id, progress_percentage)
+
+                    # Check if run has completed all epochs
+                    # completed_steps can go beyond total_steps since we stop only at a chunk boundary
+                    if run_details["completed_steps"] >= run_details["total_steps"]:
+                        scheduler.remove_run(run_id)
+                        self.db.set_run_details(
+                            run_id=run_id,
+                            status=RunStatus.COMPLETED,
+                            ended_by=RunEndedBy.EPOCH_COMPLETED,
+                        )
+                        self.logger.info(
+                            f"Run {run_id} has completed all its epochs - "
+                            f"steps {run_details['completed_steps']}/{run_details['total_steps']}"
+                        )
+                    # Check if run has completed only current epoch (hasn't reached total_steps yet)
+                    elif (
+                        new_chunks_visited == num_chunks and run_details["completed_steps"] < run_details["total_steps"]
+                    ):
+                        scheduler.reset_run(run_id)
+                        self.db.set_run_details(run_id=run_id, num_chunks_visited_curr_epoch=0)
+                        self.logger.info(f"Run {run_id} has completed epoch ({new_chunks_visited}/{num_chunks} chunks)")
+
+                # Check for failed runs and update scheduler, local state, shm
+                for worker_task in failed_tasks:
+                    run_id = worker_task["run_id"]
+                    run_error = all_run_details[run_id]["error"]
+                    if run_id in scheduler.run_ids:
+                        scheduler.remove_run(run_id)
+                        self._clear_run_from_shm(run_id)
+                        err_msg = f"Run {run_id} has failed: {run_error}"
+                        print(err_msg)
+                        self.logger.error(err_msg)
+                    self.logger.debug(f"Removed run {run_id} from scheduler")
+
+                # Process interactive control tasks (this fetches latest run states internally)
+                try:
+                    currently_scheduled_runs = list(scheduler.worker_running_current_run.values())
+                    run_states, clone_modify_tasks = self._process_interm_ic_ops_states(currently_scheduled_runs)
+                    self._process_interactive_control(
+                        run_states, clone_modify_tasks, len_train_dataset, seed, num_chunks
+                    )
+                except Exception as e:
+                    raise ControllerException(f"Error processing interactive control tasks: {e}") from e
+
+                # fetch latest run states again (post IC ops states)
+                all_run_details = self.db.get_all_runs()
+
+                # Update scheduler with active and inactive runs from IC Ops changes
+                for run_id, run_details in all_run_details.items():
+                    # add active runs to scheduler
+                    if run_details["status"] in (RunStatus.ONGOING, RunStatus.NEW) and run_id not in scheduler.run_ids:
+                        chunks_visited = all_run_details[run_id]["num_chunks_visited_curr_epoch"]
+                        start_chunk_id = all_run_details[run_id]["start_chunk_id"]
+                        scheduler.add_run(run_id, chunks_visited, start_chunk_id)
+                        self.logger.debug(f"Added run {run_id} to scheduler with {chunks_visited} chunks visited")
+                    # remove inactive runs from scheduler
+                    elif (
+                        run_details["status"] in (RunStatus.STOPPED, RunStatus.DELETED) and run_id in scheduler.run_ids
+                    ):
+                        scheduler.remove_run(run_id)
+                        self.logger.debug(f"Removed run {run_id} from scheduler")
+
+                # Get schedule from scheduler
+                schedule = scheduler.schedule()
+                run_id = schedule["run_id"]
+                worker_id = schedule["worker_id"]
+                chunk_id = schedule["chunk_id"]
+
+                # Check termination condition
+                if run_id is None and worker_id is None and chunk_id is None:
+                    self.logger.info("Scheduler indicates all runs have completed all chunks")
+                    all_done = True
+                    break
+
+                # Check if no schedule possible
+                if run_id == -1 and worker_id == -1 and chunk_id == -1:
+                    # self.logger.debug("No schedule possible - all workers busy or no available runs")
+                    time.sleep(1)
+                    continue
+
+                # Execute Schedule
+                # Create worker task
+                # self.logger.debug(f"Scheduler schedule: {schedule}")
+                self.db.set_run_details(run_id=run_id, status=RunStatus.ONGOING)
+                self.db.create_worker_task(
+                    worker_id,
+                    WorkerTask.TRAIN_VAL,
+                    TaskStatus.SCHEDULED,
+                    run_id,
+                    chunk_id,
+                    config_options={"create_model_fn": create_model_fn},
+                )
+                self.logger.debug(f"Scheduled run {run_id} on worker {worker_id} for chunk {chunk_id}")
+
+                # Small delay
+                time.sleep(1)
+
+                # Update prev_worker_tasks for next iteration (only track task_id and status)
+                prev_worker_tasks = {
+                    worker_id: {"task_id": worker_task["task_id"], "status": worker_task["status"]}
+                    for worker_id, worker_task in all_worker_tasks.items()
+                }
+
+            # set experiment task to idle
+            self.db.set_experiment_current_task(ExperimentTask.IDLE)
+            self.logger.debug(f"Set experiment task to {ExperimentTask.IDLE.value}.")
+
+        except Exception as e:
+            raise ControllerException(f"Error during run_fit: {e}") from e
+
+        # shutdown workers
+        self.worker_manager.shutdown()