from __future__ import annotations

import copy
from dataclasses import dataclass
import logging
import math
import time
from typing import Callable, Literal

import pandas as pd

from autogluon.core.models import AbstractModel, BaggedEnsembleModel
from autogluon.common.utils.resource_utils import get_resource_manager


logger = logging.getLogger(__name__)

DEFAULT_REMOTE_KWARGS = dict(max_calls=1, retry_exceptions=False, max_retries=0)


@dataclass
class ModelResources:
    """Resource allocation for a model."""

    num_gpus_for_fold_worker: int
    num_cpus_for_fold_worker: int
    num_gpus_for_model_worker: int
    num_cpus_for_model_worker: int
    total_num_cpus: int
    total_num_gpus: int


class ParallelFitManager:
    """Tracks how many resources are used when scheduling jobs in a parallel setting.

    Parallel Fit
    ---------------
    We use ray to start a model-worker with a certain number of CPUs and GPUs.
    The worker then starts new fold-workers to fit each fold of a model.
    Alternatively, the model-worker uses the given resources to perform a single fit.

    We must pass GPU resources to a model-worker if the model has `refit_folds is True`
    as the refit_folds call happens in the model-worker.

    For full parallelization, we require the following:
        - GPUs
            - refit_folds is True: `num_gpus` + `num_bag_folds` * `num_bag_sets` * `num_gpus`
            - refit_folds is False: `num_bag_folds` * `num_bag_sets` * `num_gpus`
        - CPUs:
            - model with bagging: 1 + `num_cpus` * `num_bag_folds` * `num_bag_sets`
            - model without bagging: `num_cpus`

    Parameters
    ----------
    mode: {"fit", "refit"}
        The mode to use for fitting the models.
    func: callable
        The fit function to distribute.
    func_kwargs: dict, default=None
        Additional kwargs to pass to the function.
    func_put_kwargs: dict, default=None
        Additional kwargs to pass to the function, where the values are put into the object store.
    num_cpus : int | str
        Total number of CPUs available in the cluster (or `auto`).
    num_gpus : int | str
        Total number of GPUs available in the cluster (or `auto`).
    num_splits : int | None, default=None
        Number of training splits/bags for a model. Required if mode='fit'.
    get_model_attribute_func : callable, default=None
        Function to get an attribute for a model. Required if mode='refit'.
    """
    def __init__(
        self,
        *,
        mode: Literal["fit", "refit"],
        func: Callable,
        func_kwargs: dict,
        func_put_kwargs: dict,
        num_cpus: int | str,
        num_gpus: int | str,
        num_splits: int | None = None,
        get_model_attribute_func: Callable | None = None,
        X: pd.DataFrame,
        y: pd.Series,
        problem_type="infer",
        num_classes="infer",
        total_mem: int | None | str = "auto",
        max_mem_frac: float = 0.8,
        delay_between_jobs: float = 0,
    ):
        self.X = X
        self.y = y
        self.problem_type = problem_type
        self.num_classes = num_classes
        self.delay_between_jobs = delay_between_jobs
        import ray

        self.mode = mode
        self.num_splits = num_splits
        self.get_model_attribute_func = get_model_attribute_func

        if self.mode == "fit":
            assert num_splits is not None, "num_splits must be set for mode='fit'."
        if self.mode == "refit":
            assert get_model_attribute_func is not None, "get_model_attribute_func must be set for mode='refit'."

        # Resource tracking
        if isinstance(num_cpus, str):
            num_cpus = get_resource_manager().get_cpu_count()
        if isinstance(num_gpus, str):
            num_gpus = get_resource_manager().get_gpu_count()
        self.total_num_cpus = num_cpus
        self.total_num_gpus = num_gpus
        self.available_num_cpus = num_cpus
        self.available_num_gpus = num_gpus
        # self.extra_num_cpus = self.total_num_cpus // 10  # FIXME: Maybe remove, allows for oversubscribing
        self.extra_num_cpus = 0  # Setting to 0 for now to avoid using more resources than requested by the user
        if isinstance(total_mem, str) and total_mem == "auto":
            # FIXME: when None should be infinite, implement
            total_mem = get_resource_manager().get_available_virtual_mem()
        # FIXME: Should this be passed during init or re-evaluated for each schedule call?
        self.total_mem = total_mem
        self.available_mem = self.total_mem
        self.max_mem_frac = max_mem_frac
        self.max_mem = self.total_mem * self.max_mem_frac

        # Detect max resources a job could use on some node
        self.max_cpu_resources_per_node = int(max([n["Resources"]["CPU"] for n in ray.state.nodes()]))

        # Job tracking
        self.job_refs_to_allocated_resources: dict[str, ModelResources] = {}
        self.job_refs_to_model_name: dict[str, str] = {}
        self.job_refs_to_model_memory_estimate: dict[str, int] = {}
        self.model_child_mem_estimate_cache: dict[str, int] = {}
        self.models_to_schedule: list[AbstractModel] | list[str] = []

        # Init remote function
        self.remote_func = ray.remote(**DEFAULT_REMOTE_KWARGS)(func)
        self.job_kwargs = dict()
        for key, value in func_kwargs.items():
            self.job_kwargs[key] = value
        for key, value in func_put_kwargs.items():
            self.job_kwargs[key] = ray.put(value)
        self.func_put_kwargs = func_put_kwargs

        logger.log(
            20,
            f"ParallelFitManager Details:"
            f"\n\tCPU Total: {self.total_num_cpus}"
            f"\n\tGPU Total: {self.total_num_gpus}"
            f"\n\tMem Total: {self.total_mem * 1e-9:.1f} GB"
            f"\n\t    Max Allowed: {self.max_mem * 1e-9:.1f}/{self.total_mem * 1e-9:.1f} GB (max_mem_frac={self.max_mem_frac})"
            f"\n\t    Max Allowed Per Core: {self.max_mem_per_core * 1e-9:.2f}/{self.total_mem_per_core * 1e-9:.2f} GB"
        )
    @property
    def available_num_cpus_virtual(self) -> int:
        return self.available_num_cpus + self.extra_num_cpus

    @property
    def total_num_cpus_virtual(self) -> int:
        return self.total_num_cpus + self.extra_num_cpus

    def num_children_model(self, model: AbstractModel) -> int:
        if (not isinstance(model, BaggedEnsembleModel)) or model._user_params.get("use_child_oof", False):
            return 1
        else:
            return self.num_splits

    @property
    def max_mem_per_core(self):
        return self.max_mem / self.total_num_cpus

    @property
    def total_mem_per_core(self):
        return self.total_mem / self.total_num_cpus

    # FIXME: Don't mutate models ever
    #  Mutate num_cpus user args in worker process or even better just pass num_cpus as fit kwarg
    # TODO: Don't give LightGBM/XGBoost/CatBoost more CPUs than they want, cap them at 50% of resources no matter what
    # TODO: We can lazily execute this logic. We first come up with a plan, then we schedule the workers. This allows for optimal CPU utilization.
    # FIXME v1.2: Use available memory to re-calculate per-core values, will more effectively use memory.
    # TODO: Test performance if not allowing overallocated CPUs
    def schedule_jobs(self, *, models_to_fit: list[AbstractModel] | list[str] | None = None) -> list[str]:
        """Schedule model training.

        This function must be first called with `models_to_fit is not None` and then with `models_to_fit is None`.
        Whereby the first call initializes the list of models to fit and subsequent calls schedule the remaining jobs.

        models_to_fit: list[AbstractModel] | list[str] | None, default=None
            The models that shall be fitted in a distributed manner.
        """
        import ray

        max_mem_per_core = self.max_mem_per_core

        if models_to_fit is not None:
            models_to_schedule = models_to_fit
        else:
            models_to_schedule = self.models_to_schedule

        cpus_fully_allocated = False
        num_models_delay_to_fit_all = 0

        job_refs = []
        num_models_to_schedule = len(models_to_schedule)
        models_to_schedule_tmp = []
        total_models_to_fit = 0
        for i, model in enumerate(models_to_schedule):
            model_name = model if self.mode == "refit" else model.name
            if model_name in self.model_child_mem_estimate_cache:
                model_child_memory_estimate = self.model_child_mem_estimate_cache[model_name]
            else:
                try:
                    # FIXME: DONT USE TRY/EXCEPT, this is done to handle models crashing during initialization such as KNN when `NoValidFeatures`. Instead figure this out earlier or in the worker thread
                    model_child_memory_estimate = self.get_memory_estimate_for_model_child(model=model)
                except Exception as e:
                    logger.log(20, f"Ran into exception when getting memory estimate for model, skipping model {model.name}: {e.__class__.__name__}: {e}")
                    continue
                self.model_child_mem_estimate_cache[model_name] = model_child_memory_estimate
            if model_child_memory_estimate > self.max_mem:
                logger.log(20, f"Insufficient total memory to fit model for even a single fold. Skipping {model_name}...")
                continue

            num_children = self.num_children_model(model=model)
            total_models_to_fit += num_children
            models_to_schedule_tmp.append(model)

        if total_models_to_fit == 0:
            self.models_to_schedule = []
            return []

        # Use real CPUs to avoid overallocating
        available_cpus = self.available_num_cpus
        cpus_per_model = available_cpus / total_models_to_fit
        cpus_per_model = math.floor(cpus_per_model)

        # We should always use at least this many CPUs per model
        # TODO: This isn't optimal if the user provides specific CPU requirements per model, we should take this into account
        num_cpus_per_child_floor = max(int(cpus_per_model), 1)

        models_to_schedule = models_to_schedule_tmp
        models_to_schedule_later = []
        for i, model in enumerate(models_to_schedule):
            model_name = model if self.mode == "refit" else model.name
            # TODO: refactor memory estimate logic to be one function call with the same code below
            if model_name in self.model_child_mem_estimate_cache:
                model_child_memory_estimate = self.model_child_mem_estimate_cache[model_name]
            else:
                try:
                    # FIXME: DONT USE TRY/EXCEPT, this is done to handle models crashing during initialization such as KNN when `NoValidFeatures`. Instead figure this out earlier or in the worker thread
                    model_child_memory_estimate = self.get_memory_estimate_for_model_child(model=model)
                except Exception as e:
                    logger.log(20, f"Ran into exception when getting memory estimate for model, skipping model {model.name}: {e.__class__.__name__}: {e}")
                    continue
                self.model_child_mem_estimate_cache[model_name] = model_child_memory_estimate
            if model_child_memory_estimate > self.max_mem:
                logger.log(20, f"Insufficient total memory to fit model for even a single fold. Skipping {model_name}...")
                continue
            if self.available_num_cpus_virtual < 1:
                if not cpus_fully_allocated:
                    logger.log(15, f"Delay scheduling {num_models_to_schedule - i} models: CPUs are fully allocated")
                cpus_fully_allocated = True
                models_to_schedule_later.append(model)
                continue
            num_children = self.num_children_model(model=model)
            if (num_children > self.available_num_cpus_virtual) and (self.total_num_cpus >= num_children):
                # try to wait to schedule later when all folds can be fit in parallel
                num_models_delay_to_fit_all += 1
                models_to_schedule_later.append(model)
                continue
            if num_models_delay_to_fit_all > 0:
                logger.log(
                    15,
                    f"Delay scheduling {num_models_delay_to_fit_all} models: waiting for enough CPUs to fit all folds in parallel..."
                )
                num_models_delay_to_fit_all = 0

            num_cpus_per_child_safe = model_child_memory_estimate / max_mem_per_core
            num_cpus_per_child_safe = max(math.ceil(num_cpus_per_child_safe), num_cpus_per_child_floor)

            num_cpus_avail = self.available_num_cpus_virtual

            # FIXME: Adjust num_cpus_per_child if fitting fewer folds in parallel?
            # FIXME: Not accurate, due to oversubscription, this is dangerous for memory...
            max_safe_children = math.floor(num_cpus_avail / num_cpus_per_child_safe)

            safe_children = max(min(max_safe_children, num_children), 0)

            if safe_children < num_children:
                # FIXME: Make this better, do real successive halving rather than this hack code that only works for 8 or fewer
                if safe_children >= 8:
                    safe_children = 8
                elif safe_children >= 4:
                    safe_children = 4
                elif safe_children >= 2:
                    safe_children = 2
                elif safe_children >= 1:
                    safe_children = 1
                else:
                    safe_children = 0
                    # skip model
                    pass

            if safe_children == 0:
                logger.log(15, f"Delay scheduling model {model_name}: No safe children able to be fit.")
                models_to_schedule_later.append(model)
                continue

            model_memory_estimate = self.get_memory_estimate_for_model(model=model, mem_usage_child=model_child_memory_estimate, num_children=safe_children)

            if safe_children < num_children:
                if ((num_children * model_child_memory_estimate) < self.max_mem) and (self.total_num_cpus >= num_children):
                    # try to wait to schedule later when all folds can be fit in parallel
                    logger.log(
                        15,
                        f"Delay scheduling model {model_name}: Currently can safely fit {safe_children} folds in parallel, "
                        f"waiting to be able to fit all {num_children} folds in parallel."
                    )
                    models_to_schedule_later.append(model)
                    continue
                else:
                    logger.log(
                        20,
                        f"NOTE: {model_name} is too large to ever fit all {num_children} folds in parallel. Fitting {safe_children} folds in parallel..."
                    )
                    # Will never be able to fit all children in parallel because it would use too much memory
                    # TODO: Figure out best option here, for now we train them immediately
                    #  but it may be better to deprioritize them in favor of lower memory usage models

            if isinstance(model, AbstractModel):
                # FIXME: This does in-place mutation of model, don't do this!
                # FIXME: Prefer to avoid this hack, but doing so would require some refactoring.
                #  Basically, we want to avoid passing `model` when it is initialized.
                #  Rather, it is better to pass the cls and kwargs to initialize the model downstream as late as possible.
                model = prepare_model_resources_for_fit(
                    model=model,
                    num_cpus=num_cpus_per_child_safe,
                    num_gpus=0,
                    num_parallel=safe_children,
                    num_children=num_children,
                    total_num_cpus=self.total_num_cpus,
                    total_num_gpus=self.total_num_gpus,
                )
            model_resources = self.get_resources_for_model(model=model)

            # FIXME: Keep track of memory usage estimates, use it to limit # of parallel model fits
            # FIXME: Should this logic ever trigger given the above new logic that checks for `safe_children`
            is_sufficient, reason = self.check_sufficient_resources(resources=model_resources)
            if not is_sufficient:
                if len(job_refs) + len(self.job_refs_to_allocated_resources) == 0:
                    logger.log(
                        20,
                        "DISTRIBUTED WARNING: Insufficient total resources for training a model fully distributed parallel. "
                        "Consider disabling distributed training. "
                        "Forcing to train one model anyhow, but this will lead to inefficient parallelization.",
                    )

                    # Ray's nested calls will keep blocking GPUs and thus create a deadlock if all GPUs are allocated to the model-worker and
                    # none can be used by the fold-worker.
                    if (
                        model_resources.num_gpus_for_model_worker + model_resources.num_gpus_for_fold_worker
                    ) > self.total_num_gpus:
                        raise ValueError(
                            "DISTRIBUTED ERROR: Insufficient number of GPUs to train any model, even in a non-parallel setting. "
                            "This is likely the results of requiring more GPUs than available to distribute the training. "
                            "Ray does not support freeing GPU resources for nested calls with GPUs. "
                            "Thus, we need at least twice the amount of GPUs needed to fit one model."
                        )
                else:
                    if (
                        model_resources.num_gpus_for_model_worker + model_resources.num_gpus_for_fold_worker
                    ) > self.total_num_gpus:
                        logger.log(
                            40,
                            f"DISTRIBUTED WARNING: Delay scheduling model {model_name}: "
                            "Insufficient number of GPUs to train any model, even in a non-parallel setting. "
                            "This is likely the results of requiring more GPUs than available to distribute the training. "
                            "Ray does not support freeing GPU resources for nested calls with GPUs. "
                            "Thus, we need at least twice the amount of GPUs needed to fit one model.",
                        )

                    logger.log(15, f"Delay scheduling model {model_name}: {reason}.")
                    models_to_schedule_later.append(model)
                    continue

            job_ref = self.remote_func.options(
                num_cpus=model_resources.num_cpus_for_model_worker, num_gpus=model_resources.num_gpus_for_model_worker
            ).remote(model=ray.put(model) if self.mode in ["fit"] else model, **self.job_kwargs)
            job_refs.append(job_ref)
            self.allocate_resources(job_ref=job_ref, resources=model_resources, model_name=model_name, model_memory_estimate=model_memory_estimate)

            logger.log(
                20,
                f"Scheduled {model_name}: "
                f"allocated {'' if is_sufficient else 'UP TO '}{model_resources.total_num_cpus} CPUs and {model_resources.total_num_gpus} GPUs | "
                f"{len(self.job_refs_to_allocated_resources)} jobs running"
                f"\n\t{model_resources.num_cpus_for_fold_worker if num_children != 1 else model_resources.num_cpus_for_model_worker} CPUs each for {num_children} folds, fitting {safe_children} in parallel"
                f"\n\t{self.total_num_cpus - self.available_num_cpus}/{self.total_num_cpus} Allocated CPUS"
                f"\t| {(self.total_mem - self.available_mem) * 1e-9:.1f}/{self.total_mem * 1e-9:.1f} GB Allocated Memory",
            )
            if self.delay_between_jobs > 0:
                time.sleep(self.delay_between_jobs)

        if num_models_delay_to_fit_all > 0:
            logger.log(
                15,
                f"Delay scheduling {num_models_delay_to_fit_all} models: waiting for enough CPUs to fit all folds in parallel..."
            )
            num_models_delay_to_fit_all = 0

        self.models_to_schedule = models_to_schedule_later
        return job_refs

    def check_sufficient_resources(self, *, resources: ModelResources) -> tuple[bool, str | None]:
        """Determine if there are enough resources to scheduling fitting another model."""

        # Allow for oversubscribing to 10% of the CPUs due to scheduling overhead.
        if self.available_num_cpus_virtual < resources.total_num_cpus:
            return False, "not enough CPUs free."

        # All models need at least one CPU but not all a GPU
        # Avoid scheduling a model if there are not at least 50% of the required GPUs available
        if (resources.total_num_gpus > 0) and (self.available_num_gpus < math.ceil(resources.total_num_gpus / 2)):
            return False, "not enough GPUs free."

        return True, None

    def get_resources_for_model(self, *, model: AbstractModel | str) -> ModelResources:
        if self.mode == "fit":
            # model is AbstractModel
            return self.get_resources_for_model_fit(model=model)
        elif self.mode == "refit":
            # model is str
            return self.get_resources_for_model_refit(model=model)
        else:
            raise ValueError(f"Unknown mode: {self.mode}")

    def get_memory_estimate_for_model_child(self, *, model: AbstractModel) -> int:
        X = self.X
        y = self.y

        if model.can_estimate_memory_usage_static_child():
            # This logic is fast and optimized
            # FIXME: Don't use `_user_params`
            if isinstance(model, BaggedEnsembleModel):
                hyperparameters = model.model_base._user_params
            else:
                hyperparameters = model._user_params
            mem_usage_child = model.estimate_memory_usage_static_child(
                X=X,
                y=y,
                hyperparameters=hyperparameters,
                problem_type=self.problem_type,
                num_classes=self.num_classes,
            )
            return mem_usage_child
        else:
            # This logic is over 100x slower than the static method. Not recommended.
            model_clone = copy.deepcopy(model)
            model_clone.initialize(
                X=X,
                y=y,
                problem_type=self.problem_type,
                num_classes=self.num_classes,
            )
            if isinstance(model, BaggedEnsembleModel):
                model_clone.model_base.initialize(
                    X=X,
                    y=y,
                    problem_type=self.problem_type,
                    num_classes=self.num_classes,
                )
            mem_usage_child = model_clone.estimate_memory_usage_child(
                X=X,
                y=y,
            )

            return mem_usage_child

    def get_memory_estimate_for_model(self, *, model: AbstractModel, mem_usage_child: int = None, num_children: int = None) -> int:
        if num_children is None:
            num_children = self.num_children_model(model)
        if mem_usage_child is None:
            mem_usage_child = self.get_memory_estimate_for_model_child(model=model)
        mem_usage_bag = mem_usage_child * num_children
        mem_usage_child_mb = mem_usage_child * 1e-6
        mem_usage_bag_mb = mem_usage_child_mb * num_children

        logger.log(15, f"\t{mem_usage_bag_mb:.0f} MB (per bag)\t| {mem_usage_child_mb:.0f} MB (per child)\t| {num_children} children\t| {model.name}")
        return mem_usage_bag

    def get_resources_for_model_refit(self, model: str) -> ModelResources:
        """Estimate the resources required to fit a model."""

        # FIXME: We should synchronize this with the actual values used downstream.
        #  Currently we are specifying it here and downstream separately, which gives the opportunity for them to not be matching.
        #  This should be fixed.
        # FIXME: Should this use fit_num_cpus_child or fit_num_cpus ?
        num_gpus_for_fold_worker = self.get_model_attribute_func(model=model, attribute="fit_num_gpus_child")
        num_cpus_for_fold_worker = self.get_model_attribute_func(model=model, attribute="fit_num_cpus_child")
        num_cpus_for_fold_worker = (
            num_cpus_for_fold_worker if num_cpus_for_fold_worker is not None else min(self.max_cpu_resources_per_node, self.total_num_cpus)
        )
        num_gpus_for_fold_worker = num_gpus_for_fold_worker if num_gpus_for_fold_worker is not None else 0

        num_gpus_for_model_worker = (
            1 if self.get_model_attribute_func(model=model, attribute="refit_full_requires_gpu") else 0
        )
        return ModelResources(
            num_gpus_for_fold_worker=num_gpus_for_fold_worker,
            num_cpus_for_fold_worker=num_cpus_for_fold_worker,
            num_gpus_for_model_worker=num_gpus_for_model_worker,
            num_cpus_for_model_worker=1,
            # num_cpus_for_model_worker is freed once the nested ray call is done
            total_num_cpus=num_cpus_for_fold_worker,
            total_num_gpus=num_gpus_for_model_worker + num_gpus_for_fold_worker,
        )

    # FIXME: Refactor this to be cleaner after initial merge. This code works but it is overly confusing.
    def get_resources_for_model_fit(self, *, model: AbstractModel) -> ModelResources:
        """Estimate the resources required to fit a model."""

        if "num_cpus" not in getattr(model, "model_base", model)._user_params_aux:
            logger.warning(
                f"DISTRIBUTED WARNING: Model {model.name} does not specify the number of resources to use! "
                "Assuming that the model will use all available node resources, which can heavily impact the performance of distributed training."
            )

        # Fallback if required information are not given at this point, we can only assume that the model will use all available CPU resources of a node.
        num_cpus_for_fold_worker = getattr(model, "model_base", model)._user_params_aux.get(
            "num_cpus", self.max_cpu_resources_per_node
        )
        # As we have no information about if the model needs GPUs, we assume that it does not need any.
        num_gpus_for_fold_worker = getattr(model, "model_base", model)._user_params_aux.get("num_gpus", 0)

        if (not isinstance(model, BaggedEnsembleModel)) or model._user_params.get("use_child_oof", False):
            # Only one fold is fit, so we need to use all fit resources for the model-worker
            num_cpus_for_model_worker = num_cpus_for_fold_worker
            num_gpus_for_model_worker = num_gpus_for_fold_worker
            num_cpus_for_fold_worker = 0
            num_gpus_for_fold_worker = 0
            total_num_cpus = num_cpus_for_model_worker
        else:
            # If refit_folds is True, we need to pass GPU resources to the model-worker
            num_gpus_for_model_worker = (
                num_gpus_for_fold_worker
                if ((num_gpus_for_fold_worker > 0) and model._user_params.get("refit_folds", False))
                else 0
            )
            num_cpus_for_model_worker = 1

            # num_cpus_for_model_worker is freed once the nested ray call is done
            total_num_cpus = model._user_params_aux["num_cpus"]

        return ModelResources(
            num_gpus_for_fold_worker=num_gpus_for_fold_worker,
            num_cpus_for_fold_worker=num_cpus_for_fold_worker,
            num_gpus_for_model_worker=num_gpus_for_model_worker,
            num_cpus_for_model_worker=num_cpus_for_model_worker,
            total_num_cpus=total_num_cpus,
            total_num_gpus=num_gpus_for_model_worker + num_gpus_for_fold_worker * self.num_splits,
        )

    def allocate_resources(self, *, job_ref: str, resources: ModelResources, model_memory_estimate: int, model_name: str = None) -> None:
        """Allocate resources for a model fit."""

        self.available_num_cpus -= resources.total_num_cpus
        self.available_num_gpus -= resources.total_num_gpus
        self.available_mem -= model_memory_estimate
        self.job_refs_to_allocated_resources[job_ref] = resources
        self.job_refs_to_model_name[job_ref] = model_name
        self.job_refs_to_model_memory_estimate[job_ref] = model_memory_estimate

    def deallocate_resources(self, *, job_ref: str) -> None:
        """Deallocate resources for a model fit."""

        resources = self.job_refs_to_allocated_resources.pop(job_ref)
        model_name = self.job_refs_to_model_name.pop(job_ref)
        self.available_num_cpus += resources.total_num_cpus
        self.available_num_gpus += resources.total_num_gpus
        model_memory_estimate = self.job_refs_to_model_memory_estimate.pop(job_ref)
        self.available_mem += model_memory_estimate
        self.model_child_mem_estimate_cache.pop(model_name)

    def clean_unfinished_job_refs(self, *, unfinished_job_refs: list[str] | None = None):
        import ray

        # TODO: determine how to suppress error messages from cancelling jobs.
        if unfinished_job_refs is not None and len(unfinished_job_refs) > 0:
            model_names = [self.job_refs_to_model_name[f] for f in unfinished_job_refs]
            logger.log(20, f"Cancelling {len(model_names)} jobs: {model_names}")
            for f in unfinished_job_refs:
                ray.cancel(f, force=True)

    def clean_job_state(self, *, unfinished_job_refs: list[str] | None = None) -> None:
        """Clean up state of manager."""
        self.job_refs_to_allocated_resources = {}
        self.job_refs_to_model_name = {}
        self.job_refs_to_model_memory_estimate = {}
        self.model_child_mem_estimate_cache = {}
        self.models_to_schedule = []
        self.available_num_cpus = self.total_num_cpus
        self.available_num_gpus = self.total_num_gpus
        self.available_mem = self.total_mem
        self.clean_unfinished_job_refs(unfinished_job_refs=unfinished_job_refs)

    def clean_up_ray(self, *, unfinished_job_refs: list[str] | None = None) -> None:
        """Try to clean up ray object store."""
        import ray

        self.clean_unfinished_job_refs(unfinished_job_refs=unfinished_job_refs)

        ray.internal.free(object_refs=[self.job_kwargs[key] for key in self.func_put_kwargs])
        for key in self.func_put_kwargs:
            del self.job_kwargs[key]


# TODO: This logic is a hack. We shouldn't be editing models in-place. Refactor this as a fast-follow. Currently it works, but it is not ideal long term.
def prepare_model_resources_for_fit(
    *,
    model: AbstractModel,
    total_num_cpus: int,
    total_num_gpus: int,
    num_cpus: int = 1,
    num_gpus: float = 0,
    num_cpus_worker: int = 1,
    num_gpus_worker: float = 0,
    num_parallel: int = None,
    num_children: int = 1,
) -> AbstractModel:
    """Allocate each model resources for fitting. (This is currently an in-place operation!)

    We allocate resources by setting the _user_params_aux of a model.

    """
    if num_parallel is None:
        num_parallel = num_children
    assert num_parallel <= num_children

    if num_cpus == 0:
        # Non-bagged mode
        num_cpus = num_cpus_worker
        num_gpus = num_gpus_worker
    else:
        num_cpus_parent = num_cpus * num_parallel
        num_gpus_parent = num_gpus * num_parallel

        model_aux = model._user_params_aux
        if "num_cpus" not in model_aux:
            model_aux["num_cpus"] = num_cpus_parent
        if "num_gpus" not in model_aux:
            model_aux["num_gpus"] = num_gpus_parent

    upa = getattr(model, "model_base", model)._user_params_aux
    if "num_cpus" not in upa:
        getattr(model, "model_base", model)._user_params_aux["num_cpus"] = num_cpus
    if "num_gpus" not in upa:
        getattr(model, "model_base", model)._user_params_aux["num_gpus"] = num_gpus

    return model
