rapidfireai 0.10.2rc5__py3-none-any.whl → 0.11.1rc1__py3-none-any.whl

rapidfireai/frontend/proxy_middleware.py

@@ -25,7 +25,7 @@ class UserProxyManager:
         self.default_proxy = {
             'main_proxy_target': 'http://127.0.0.1:5002/',
             'static_proxy_target': 'http://127.0.0.1:5002/',
-            'dispatcher_proxy_target': 'http://127.0.0.1:8080/',
+            'dispatcher_proxy_target': 'http://127.0.0.1:8081/',
         }
     
     def get_user_proxy(self, user_id: str) -> Dict[str, str]:

rapidfireai/ml/callbacks.py

@@ -1,14 +1,9 @@
-from typing import Callable, Dict, List, Optional
+from collections.abc import Callable
 
 import torch
 from datasets import Dataset
 from tqdm import tqdm
-from transformers import (
-    TrainerCallback,
-    TrainerControl,
-    TrainerState,
-    TrainingArguments,
-)
+from transformers import TrainerCallback, TrainerControl, TrainerState, TrainingArguments
 from transformers.trainer_utils import IntervalStrategy, SaveStrategy
 
 
@@ -17,10 +12,10 @@ class GenerationMetricsCallback(TrainerCallback):
         self,
         tokenizer,
         eval_dataset: Dataset,
-        generation_config: Optional[Dict] = None,
+        generation_config: dict | None = None,
         compute_metrics: Callable = None,
         batch_size: int = 8,
-        mlflow_manager=None,
+        metric_logger=None,
         mlflow_run_id: str = None,
         completed_steps: int = 0,
     ):
@@ -36,7 +31,7 @@ class GenerationMetricsCallback(TrainerCallback):
             "pad_token_id": tokenizer.pad_token_id,
             "eos_token_id": tokenizer.eos_token_id,
         }
-        self.mlflow_manager = mlflow_manager
+        self.metric_logger = metric_logger
         self.mlflow_run_id = mlflow_run_id
         self.completed_steps = completed_steps
 
@@ -63,8 +58,8 @@ class GenerationMetricsCallback(TrainerCallback):
             state.log_history.append(metrics)
 
         for key, value in metrics.items():
-            if self.mlflow_manager:
-                self.mlflow_manager.log_metric(
+            if self.metric_logger:
+                self.metric_logger.log_metric(
                     self.mlflow_run_id,
                     key,
                     value,
@@ -72,43 +67,69 @@ class GenerationMetricsCallback(TrainerCallback):
                 )
 
     def _prepare_data(self, eval_dataset: Dataset) -> tuple:
-        """Prepare batch data for generation"""
+        """Prepare batch data for generation with defensive validation"""
         input_texts = []
         references = []
 
         for item in eval_dataset:
-            if isinstance(item, dict):
-                if "input" in item and "output" in item:
-                    input_text = item["input"]
-                    reference = item["output"]
-                elif "prompt" in item and "completion" in item:
-                    input_text = item["prompt"]
-                    reference = item["completion"][-1]["content"]
-                    input_text = self.tokenizer.apply_chat_template(
-                        input_text, tokenize=False
-                    )
-                else:
-                    continue
-
-                input_texts.append(input_text)
-                references.append(reference)
+            if not isinstance(item, dict):
+                continue
+
+            input_text = None
+            reference = None
+
+            # Support multiple field name patterns
+            if "input" in item and "output" in item:
+                input_text = item["input"]
+                reference = item["output"]
+            elif "prompt" in item and "completion" in item:
+                input_text = item["prompt"]
+                reference = item["completion"][-1]["content"]
+                input_text = self.tokenizer.apply_chat_template(input_text, tokenize=False)
+            elif "text" in item:
+                # SFT format - use text as input, response as reference
+                input_text = item["text"]
+                reference = item.get("response", item.get("instruction", item["text"]))
+            elif "instruction" in item and "response" in item:
+                # Direct instruction/response format
+                input_text = item["instruction"]
+                reference = item["response"]
+
+            # Validate non-empty strings
+            if input_text and isinstance(input_text, str) and input_text.strip():
+                if reference and isinstance(reference, str) and reference.strip():
+                    input_texts.append(input_text.strip())
+                    references.append(reference.strip())
+
+        # Return safe empty values to prevent downstream errors
+        if not input_texts:
+            return [], []
 
         return input_texts, references
 
-    def _generate_batch(self, model, input_texts: List[str]) -> List[str]:
-        """Generate text for a batch of inputs"""
-        # Tokenize batch
-        inputs = self.tokenizer(
-            input_texts,
-            return_tensors="pt",
-            padding=True,
-            truncation=True,
-            max_length=512,  # Adjust based on your model's context length
-        ).to(model.device)
+    def _generate_batch(self, model, input_texts: list[str]) -> torch.Tensor:
+        """Generate text for a batch of inputs with defensive validation"""
+        # Defensive validation for empty inputs
+        if not input_texts:
+            return torch.empty((0, 0), dtype=torch.long).to(model.device)
 
-        return inputs["input_ids"]
-
-    def _compute_generation_metrics(self, model, step: int) -> Dict[str, float]:
+        try:
+            # Tokenize batch
+            inputs = self.tokenizer(
+                input_texts,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=512,  # Adjust based on your model's context length
+            ).to(model.device)
+
+            return inputs["input_ids"]
+        except Exception as e:
+            # Log error and return empty tensor to prevent crash
+            print(f"Warning: Tokenization error in generation callback: {e}")
+            return torch.empty((0, 0), dtype=torch.long).to(model.device)
+
+    def _compute_generation_metrics(self, model, step: int) -> dict[str, float]:
         """Generate text and compute BLEU/ROUGE metrics with batch processing"""
         model.eval()
 
@@ -121,16 +142,24 @@ class GenerationMetricsCallback(TrainerCallback):
 
         # Process in batches
         input_texts, batch_references = self._prepare_data(self.eval_dataset)
+
+        # Early return if no valid data
+        if not input_texts:
+            print("Warning: No valid eval data for generation metrics")
+            return {}
+
         input_ids = self._generate_batch(model, input_texts)
+
+        # Check for empty generation batch
+        if input_ids.numel() == 0:
+            print("Warning: Empty input_ids from tokenization")
+            return {}
+
         with torch.no_grad():
-            for i in tqdm(
-                range(0, len(indices), self.batch_size), desc="Generating for metrics"
-            ):
+            for i in tqdm(range(0, len(indices), self.batch_size), desc="Generating for metrics"):
                 input_ids_batch = input_ids[i : i + self.batch_size]
                 with torch.inference_mode(), torch.cuda.amp.autocast():
-                    outputs_batch = model.generate(
-                        input_ids_batch, **self.generation_config
-                    )
+                    outputs_batch = model.generate(input_ids_batch, **self.generation_config)
                 generated_texts = self.tokenizer.batch_decode(
                     outputs_batch[:, input_ids_batch.shape[1] :],
                     skip_special_tokens=True,
@@ -155,18 +184,18 @@ class GenerationMetricsCallback(TrainerCallback):
 
 
 class MLflowLoggingCallback(TrainerCallback):
-    """Callback for logging metrics to MLflow during training"""
+    """Callback for logging metrics to tracking backend during training"""
 
     def __init__(
         self,
-        mlflow_manager,
+        metric_logger,
         mlflow_run_id: str,
         excluded_keys: list = None,
         completed_steps: int = 0,
         chunk_id: int = 0,
         num_epochs_completed: int = 0,
     ):
-        self.mlflow_manager = mlflow_manager
+        self.metric_logger = metric_logger
         self.mlflow_run_id = mlflow_run_id
         self.completed_steps = completed_steps
         self.excluded_keys = excluded_keys or [
@@ -189,22 +218,22 @@ class MLflowLoggingCallback(TrainerCallback):
             for key, value in logs.items():
                 if isinstance(value, (int, float)) and key not in self.excluded_keys:
                     try:
-                        self.mlflow_manager.log_metric(
+                        self.metric_logger.log_metric(
                             self.mlflow_run_id,
                             key,
                             value,
                             step=self.completed_steps + state.global_step,
                         )
                     except Exception as e:
-                        print(f"Warning: Failed to log metric {key} to MLflow: {e}")
+                        print(f"Warning: Failed to log metric {key} to tracking backend: {e}")
             if "eval_loss" not in logs and "train_runtime" not in logs:
-                self.mlflow_manager.log_metric(
+                self.metric_logger.log_metric(
                     self.mlflow_run_id,
                     "chunk number",
                     self.chunk_id,
                     step=self.completed_steps + state.global_step,
                 )
-                self.mlflow_manager.log_metric(
+                self.metric_logger.log_metric(
                     self.mlflow_run_id,
                     "num_epochs_completed",
                     self.num_epochs_completed,
@@ -217,7 +246,7 @@ class LogLevelCallback(TrainerCallback):
     A [`TrainerCallback`] that handles the default flow of the training loop for logs, evaluation and checkpoints.
     """
 
-    def __init__(self, global_step_args: Dict):
+    def __init__(self, global_step_args: dict):
         self.eval_first_step = global_step_args.get("eval_first_step", 0)
         self.actual_steps = global_step_args.get("actual_steps", 0)
         self.log_first_step = global_step_args.get("log_first_step", 0)
@@ -275,10 +304,7 @@ class LogLevelCallback(TrainerCallback):
             control.should_log = True
 
         # Evaluate
-        if (
-            args.eval_strategy == IntervalStrategy.EPOCH
-            and args.eval_delay <= state.epoch
-        ):
+        if args.eval_strategy == IntervalStrategy.EPOCH and args.eval_delay <= state.epoch:
             control.should_evaluate = True
 
         # Save

rapidfireai/ml/trainer.py

@@ -1,4 +1,3 @@
-import logging
 import math
 import os
 
@@ -7,11 +6,7 @@ from peft import LoraConfig, get_peft_model_state_dict, set_peft_model_state_dic
 from transformers.utils.logging import set_verbosity_error
 from trl import DPOConfig, DPOTrainer, GRPOConfig, GRPOTrainer, SFTConfig, SFTTrainer
 
-from rapidfireai.ml.callbacks import (
-    GenerationMetricsCallback,
-    MLflowLoggingCallback,
-    LogLevelCallback,
-)
+from rapidfireai.ml.callbacks import GenerationMetricsCallback, LogLevelCallback, MLflowLoggingCallback
 from rapidfireai.ml.checkpoint_utils import (
     ensure_gradient_compatibility,
     load_checkpoint_from_disk,
@@ -34,7 +29,7 @@ def create_trainer_instance(
     trainer_config: TrainerConfig,
     shm_manager: SharedMemoryManager,
     use_shared_memory: bool = False,
-    mlflow_manager=None,
+    metric_logger=None,
     chunk_id: int = 0,
 ) -> tuple[SFTTrainer | DPOTrainer | GRPOTrainer | None, str]:
     """
@@ -51,21 +46,15 @@ def create_trainer_instance(
     compute_metrics = additional_trainer_kwargs.get("compute_metrics", None)
 
     # Configure training arguments
-    training_args, global_step_args = _configure_training_args(
-        training_args, trainer_config
-    )
+    training_args, global_step_args = _configure_training_args(training_args, trainer_config)
     trainer_config_obj = _create_trainer_config_object(trainer_type, training_args)
     # check if peft params is empty dict
     is_peft = bool(config_leaf.get("peft_params"))
     # Load model and tokenizer
     if use_shared_memory:
-        model_instance, tokenizer = load_checkpoint_from_shared_memory(
-            trainer_config, shm_manager, is_peft=is_peft
-        )
+        model_instance, tokenizer = load_checkpoint_from_shared_memory(trainer_config, shm_manager, is_peft=is_peft)
     else:
-        model_instance, tokenizer = load_checkpoint_from_disk(
-            trainer_config, is_peft=is_peft
-        )
+        model_instance, tokenizer = load_checkpoint_from_disk(trainer_config, is_peft=is_peft)
     # add model name to model config
     config_leaf["model_name"] = model_instance.config._name_or_path
 
@@ -84,30 +73,26 @@ def create_trainer_instance(
 
     model_instance = model_instance.to(device)
 
-    trainer_kwargs, formatting_func, additional_trainer_kwargs = (
-        _prepare_trainer_kwargs(
-            model_instance,
-            trainer_config_obj,
-            tokenizer,
-            trainer_config,
-            additional_trainer_kwargs,
-            ref_model_instance,
-            config_leaf,
-        )
+    trainer_kwargs, formatting_func, additional_trainer_kwargs = _prepare_trainer_kwargs(
+        model_instance,
+        trainer_config_obj,
+        tokenizer,
+        trainer_config,
+        additional_trainer_kwargs,
+        ref_model_instance,
+        config_leaf,
     )
 
-    callbacks, additional_trainer_kwargs = (
-        _setup_callbacks(  # FIXME: avoid returning additional_trainer_kwargs
-            mlflow_manager,
-            trainer_config,
-            chunk_id,
-            compute_metrics,
-            additional_trainer_kwargs,
-            tokenizer,
-            training_args,
-            formatting_func,
-            global_step_args,
-        )
+    callbacks, additional_trainer_kwargs = _setup_callbacks(  # FIXME: avoid returning additional_trainer_kwargs
+        metric_logger,
+        trainer_config,
+        chunk_id,
+        compute_metrics,
+        additional_trainer_kwargs,
+        tokenizer,
+        training_args,
+        formatting_func,
+        global_step_args,
     )
 
     if callbacks:
@@ -116,29 +101,22 @@ def create_trainer_instance(
     trainer_kwargs.update(additional_trainer_kwargs)
     trainer_kwargs = {k: v for k, v in trainer_kwargs.items() if v is not None}
 
-    trainer = _create_trainer_by_type(
-        trainer_type, trainer_kwargs, trainer_config, use_shared_memory, shm_manager
-    )
+    trainer = _create_trainer_by_type(trainer_type, trainer_kwargs, trainer_config, use_shared_memory, shm_manager)
     return trainer, config_leaf["model_name"]
 
 
-def _configure_training_args(
-    training_args: dict, trainer_config: TrainerConfig
-) -> dict:
+def _configure_training_args(training_args: dict, trainer_config: TrainerConfig) -> dict:
     """Configure training arguments with default values."""
     completed_steps = trainer_config.completed_steps
     per_device_train_batch_size = training_args.get("per_device_train_batch_size", 1)
     gradient_accumulation_steps = training_args.get("gradient_accumulation_steps", 1)
-    len_dataloader = math.ceil(
-        trainer_config.train_dataset.num_rows / per_device_train_batch_size
-    )
+    len_dataloader = math.ceil(trainer_config.train_dataset.num_rows / per_device_train_batch_size)
     steps_per_epoch = max(
-        len_dataloader // gradient_accumulation_steps
-        + int(len_dataloader % gradient_accumulation_steps > 0),
+        len_dataloader // gradient_accumulation_steps + int(len_dataloader % gradient_accumulation_steps > 0),
         1,
     )
 
-    if trainer_config.config_leaf.get("trainer_type","SFT") == "GRPO":
+    if trainer_config.config_leaf.get("trainer_type", "SFT") == "GRPO":
         num_generations = training_args.get("num_generations", 8)
         steps_per_epoch = (num_generations * trainer_config.train_dataset.num_rows) // (
             gradient_accumulation_steps * per_device_train_batch_size
@@ -215,10 +193,7 @@ def _setup_reference_model(
         if model_adapter_name is not None and ref_adapter_name is not None:
             if use_shared_memory:
                 peft_config = LoraConfig(**config_leaf["peft_params"])
-                if (
-                    trainer_config.completed_steps == 0
-                    and trainer_config.warm_started_from is None
-                ):
+                if trainer_config.completed_steps == 0 and trainer_config.warm_started_from is None:
                     reference_state_dict = get_peft_model_state_dict(model_instance)
                     reference_state_dict = move_tensors_to_cpu(reference_state_dict)
                     shm_manager.save_model_object(
@@ -230,14 +205,10 @@ def _setup_reference_model(
                     reference_state_dict = shm_manager.load_model_object(
                         trainer_config.run_id, SHMObjectType.REF_STATE_DICT
                     )
-                    reference_state_dict = move_tensors_to_device(
-                        reference_state_dict, device
-                    )
+                    reference_state_dict = move_tensors_to_device(reference_state_dict, device)
                 model_instance.add_adapter(ref_adapter_name, peft_config)
                 model_instance.set_adapter(ref_adapter_name)
-                set_peft_model_state_dict(
-                    model_instance, reference_state_dict, adapter_name=ref_adapter_name
-                )
+                set_peft_model_state_dict(model_instance, reference_state_dict, adapter_name=ref_adapter_name)
                 model_instance.set_adapter(model_adapter_name)
             else:
                 base_run_path = DataPath.base_run_path(trainer_config.run_id)
@@ -289,9 +260,7 @@ def _prepare_trainer_kwargs(
 
     if additional_trainer_kwargs.get("formatting_func") is not None:
         formatting_func = additional_trainer_kwargs.get("formatting_func")
-        train_dataset = train_dataset.map(
-            formatting_func
-        )  # FIXME: add try exception with batched/unbatched
+        train_dataset = train_dataset.map(formatting_func)  # FIXME: add try exception with batched/unbatched
         if eval_dataset is not None:
             eval_dataset = eval_dataset.map(formatting_func)
         additional_trainer_kwargs_copy = additional_trainer_kwargs.copy()
@@ -314,7 +283,7 @@ def _prepare_trainer_kwargs(
 
 
 def _setup_callbacks(
-    mlflow_manager,
+    metric_logger,
     trainer_config,
     chunk_id,
     compute_metrics,
@@ -327,9 +296,9 @@ def _setup_callbacks(
     """Setup callbacks for the trainer."""
     callbacks = []
 
-    if mlflow_manager is not None and trainer_config.mlflow_run_id is not None:
+    if metric_logger is not None and trainer_config.mlflow_run_id is not None:
         mlflow_callback = MLflowLoggingCallback(
-            mlflow_manager=mlflow_manager,
+            metric_logger=metric_logger,
             mlflow_run_id=trainer_config.mlflow_run_id,
             completed_steps=trainer_config.completed_steps,
             chunk_id=chunk_id,
@@ -337,10 +306,7 @@ def _setup_callbacks(
         )
         callbacks.append(mlflow_callback)
 
-    if (
-        compute_metrics is not None
-        and additional_trainer_kwargs.get("generation_config") is not None
-    ):
+    if compute_metrics is not None and additional_trainer_kwargs.get("generation_config") is not None:
         compute_metrics_function = compute_metrics
         if formatting_func is not None:
             formatted_eval_dataset = trainer_config.eval_dataset.map(formatting_func)
@@ -353,7 +319,7 @@ def _setup_callbacks(
             generation_config=additional_trainer_kwargs.get("generation_config"),
             compute_metrics=compute_metrics_function,
             batch_size=training_args.get("per_device_eval_batch_size"),
-            mlflow_manager=mlflow_manager,
+            metric_logger=metric_logger,
             mlflow_run_id=trainer_config.mlflow_run_id,
             completed_steps=trainer_config.completed_steps,
         )
@@ -365,15 +331,11 @@ def _setup_callbacks(
     return callbacks, additional_trainer_kwargs
 
 
-def _create_trainer_by_type(
-    trainer_type, trainer_kwargs, trainer_config, use_shared_memory, shm_manager
-):
+def _create_trainer_by_type(trainer_type, trainer_kwargs, trainer_config, use_shared_memory, shm_manager):
     """Create trainer instance based on type with proper state restoration."""
     if trainer_type == "SFT":
         dummy_trainer = SFTTrainer(**trainer_kwargs)
-        dummy_trainer.create_optimizer_and_scheduler(
-            num_training_steps=trainer_config.total_steps
-        )
+        dummy_trainer.create_optimizer_and_scheduler(num_training_steps=trainer_config.total_steps)
         trainer = SFTTrainer(
             **trainer_kwargs,
             optimizers=(dummy_trainer.optimizer, dummy_trainer.lr_scheduler),
@@ -382,9 +344,7 @@ def _create_trainer_by_type(
 
     elif trainer_type == "DPO":
         dummy_trainer = DPOTrainer(**trainer_kwargs)
-        dummy_trainer.create_optimizer_and_scheduler(
-            num_training_steps=trainer_config.total_steps
-        )
+        dummy_trainer.create_optimizer_and_scheduler(num_training_steps=trainer_config.total_steps)
         trainer = DPOTrainer(
             **trainer_kwargs,
             optimizers=(dummy_trainer.optimizer, dummy_trainer.lr_scheduler),
@@ -393,9 +353,7 @@ def _create_trainer_by_type(
 
     elif trainer_type == "GRPO":
         dummy_trainer = GRPOTrainer(**trainer_kwargs)
-        dummy_trainer.create_optimizer_and_scheduler(
-            num_training_steps=trainer_config.total_steps
-        )
+        dummy_trainer.create_optimizer_and_scheduler(num_training_steps=trainer_config.total_steps)
         trainer = GRPOTrainer(
             **trainer_kwargs,
             optimizers=(dummy_trainer.optimizer, dummy_trainer.lr_scheduler),
@@ -406,9 +364,7 @@ def _create_trainer_by_type(
 
     if trainer_config.completed_steps > 0:
         if use_shared_memory:
-            trainer = restore_trainer_from_shared_memory(
-                trainer, trainer_config, shm_manager
-            )
+            trainer = restore_trainer_from_shared_memory(trainer, trainer_config, shm_manager)
         else:
             trainer = restore_trainer_from_disk(trainer, trainer_config)