Feature/fit on ray

pyproject.toml

@@ -95,7 +95,7 @@ colab_evals = [
     # Statistical Analysis
 
     # Data Manipulation & Display
-    "unstructured>=0.18.15",   
+    "unstructured>=0.18.15",
 ]
 colab_fit = [
     "pandas>=2.0.0,<2.3.0",  # Colab compatibility (2.2.2) and cudf constraint
@@ -109,9 +109,9 @@ colab_fit = [
     "evaluate>=0.4.5",
     "rouge-score>=0.1.2",
     "sentencepiece>=0.2.1",
+    "ray[default]>=2.46.0",  # Ray for distributed training
     "mlflow>=3.2.0",
     "requests>=2.32.0",  # Relaxed for Colab (2.32.4)
-    "loguru>=0.7.3",
     "ipython>=7.34.0",  # Colab compatibility (7.34.0)
     "jupyter>=1.1.1",
     "ipywidgets>=7.3.4,<9.0.0",  # Support both v7 (Colab) and v8 (Jupyter)
@@ -178,7 +178,6 @@ local_evals = [
     "sentencepiece>=0.2.1",
     "mlflow>=3.2.0",
     "requests>=2.32.0",  # Relaxed for Colab (2.32.4)
-    "loguru>=0.7.3",
     "ipython>=7.34.0",  # Colab compatibility (7.34.0)
     "jupyter>=1.1.1",
     "flashinfer-python==0.2.5",
@@ -245,7 +244,6 @@ local_fit = [
     "dill>=0.3.8",
     "mlflow>=3.2.0",
     "requests>=2.32.0",  # Relaxed for Colab (2.32.4)
-    "loguru>=0.7.3",
     "ipython>=7.34.0",  # Colab compatibility (7.34.0)
     "jupyter>=1.1.1",
     "trackio",

rapidfireai/automl/automl_utils.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 from rapidfireai.automl.base import AutoMLAlgorithm
-from rapidfireai.fit.utils.exceptions import AutoMLException
+from rapidfireai.utils.exceptions import AutoMLException
 
 
 def _is_valid_reranker_top_n_vs_k(pipeline: Any) -> bool:

rapidfireai/automl/base.py

@@ -4,7 +4,7 @@
 from typing import Any
 
 from rapidfireai.automl.datatypes import List
-from rapidfireai.fit.utils.exceptions import AutoMLException
+from rapidfireai.utils.exceptions import AutoMLException
 
 
 class AutoMLAlgorithm(ABC):

rapidfireai/automl/grid_search.py

@@ -7,7 +7,7 @@
 from rapidfireai.automl.base import AutoMLAlgorithm
 from rapidfireai.automl.datatypes import List
 from rapidfireai.automl.automl_utils import filter_evals_runs_valid_reranker
-from rapidfireai.fit.utils.exceptions import AutoMLException
+from rapidfireai.utils.exceptions import AutoMLException
 
 
 def recursive_expand_gridsearch(item: Any):

rapidfireai/automl/model_config.py

@@ -6,7 +6,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Type, get_type_hints
+from typing import Any, get_type_hints
 
 from rapidfireai.automl.datatypes import List, Range
 
@@ -207,7 +207,7 @@ def get_engine_kwargs(self) -> dict[str, Any]:
             pass
 
 
-def _create_rf_class_evals(base_class: Type, class_name: str):
+def _create_rf_class_evals(base_class: type, class_name: str):
     """Creating a RF class for evals that dynamically inherits all constructor parameters and supports singleton, list, and Range values."""
     if not inspect.isclass(base_class):
         raise ValueError(f"base_class must be a class, got {type(base_class)}")

rapidfireai/automl/random_search.py

@@ -1,14 +1,14 @@
 """Random search implementation for AutoML hyperparameter optimization."""
 
+import hashlib
 import json
 import random
-import hashlib
 from typing import Any
 
 from rapidfireai.automl.base import AutoMLAlgorithm
 from rapidfireai.automl.datatypes import List, Range
 from rapidfireai.automl.automl_utils import _is_valid_reranker_top_n_vs_k
-from rapidfireai.fit.utils.exceptions import AutoMLException
+from rapidfireai.utils.exceptions import AutoMLException
 
 
 def encode_payload(payload: dict[str, Any]) -> str:
@@ -73,16 +73,12 @@ def _get_runs_fit(self, seed: int) -> list[dict[str, Any]]:
                 selected_peft_config = config.peft_config
 
             peft_params = (
-                {}
-                if selected_peft_config is None
-                else recursive_expand_randomsearch(selected_peft_config._user_params)
+                {} if selected_peft_config is None else recursive_expand_randomsearch(selected_peft_config._user_params)
             )
 
             # Sample other parameters
             training_params = (
-                {}
-                if config.training_args is None
-                else recursive_expand_randomsearch(config.training_args._user_params)
+                {} if config.training_args is None else recursive_expand_randomsearch(config.training_args._user_params)
             )
 
             model_kwargs = (
@@ -202,11 +198,7 @@ def _get_runs_evals(self, seed: int) -> list[dict[str, Any]]:
 
             for pipeline in pipelines:
                 # Sample model config parameters
-                pipeline_instances = (
-                    [{}]
-                    if pipeline is None
-                    else [recursive_expand_randomsearch(pipeline)]
-                )
+                pipeline_instances = [{}] if pipeline is None else [recursive_expand_randomsearch(pipeline)]
 
                 additional_kwargs = {
                     k: v
@@ -217,9 +209,7 @@ def _get_runs_evals(self, seed: int) -> list[dict[str, Any]]:
                     and v is not None
                 }
                 additional_kwargs_instances = (
-                    [{}]
-                    if not additional_kwargs
-                    else [recursive_expand_randomsearch(additional_kwargs)]
+                    [{}] if not additional_kwargs else [recursive_expand_randomsearch(additional_kwargs)]
                 )
 
                 # Generate random search combinations