arbor-ai 0.2.1__py3-none-any.whl → 0.2.2__py3-none-any.whl

arbor/server/services/scripts/grpo_training.py

@@ -5,24 +5,19 @@
 
 import argparse
 import json
-import os
 import random
-import shutil
 import signal
 import sys
 import threading
 import time
-from functools import lru_cache
 from typing import Any, List, Optional, Union
 
 import torch
 import trl.extras.vllm_client
 import zmq
-from accelerate import Accelerator
 from accelerate.utils import broadcast_object_list, gather, gather_object
 from datasets import Dataset, IterableDataset, load_dataset
-from peft import AutoPeftModelForCausalLM, LoraConfig, PeftConfig  # type: ignore
-from torch.utils.data import Dataset
+from peft import LoraConfig, PeftConfig
 from transformers import (
     PreTrainedModel,
     PreTrainedTokenizerBase,
@@ -38,28 +33,19 @@ from arbor.server.services.comms.comms import (
     ArborServerCommsHandler,
 )
 from arbor.server.services.inference.vllm_client import VLLMClient
+from arbor.server.services.scripts.utils.callbacks import WeightUpdateCallback
+from arbor.server.services.scripts.utils.comms_monitors import CommandMonitor
+from arbor.server.services.scripts.utils.dataset import BlockingRotatingQueueDataset
+from arbor.server.services.scripts.utils.ingestion_monitor import IngestionMonitor
 
 trl.extras.vllm_client.VLLMClient = VLLMClient
 
-if is_wandb_available():
-    import wandb
-
-last_step_time = None
-last_queue_pop_time = None
-
-
-def time_since_last_step():
-    global last_step_time
-    if last_step_time is None:
-        return float("inf")
-    return time.time() - last_step_time
+from arbor.server.utils.logging import get_logger
 
+logger = get_logger(__name__)
 
-def get_time_since_last_queue_pop():
-    global last_queue_pop_time
-    if last_queue_pop_time is None:
-        return float("inf")
-    return time.time() - last_queue_pop_time
+if is_wandb_available():
+    import wandb
 
 
 class ArborGRPOTrainer(GRPOTrainer):
@@ -77,10 +63,7 @@ class ArborGRPOTrainer(GRPOTrainer):
         ] = (None, None),
         peft_config: Optional["PeftConfig"] = None,
         comms_handler: Optional[ArborScriptCommsHandler] = None,
-        lora: Optional[bool] = False,
-        # We do nothing with max_context_length right now
         vllm_group_port: Optional[int] = None,
-        max_context_length: Optional[int] = None,
         **kwargs,
     ):
         super().__init__(
@@ -103,7 +86,7 @@ class ArborGRPOTrainer(GRPOTrainer):
         args.use_vllm = True
         self.use_vllm = True
         if self.accelerator.is_main_process:
-            print(
+            logger.info(
                 f"Initializing vLLM client with server port {args.vllm_server_port} and group port {vllm_group_port}"
             )
             self.vllm_client = VLLMClient(
@@ -185,13 +168,15 @@ class ArborGRPOTrainer(GRPOTrainer):
 
         if self.max_prompt_length is not None:
             if prompt_ids.shape[1] > self.max_prompt_length:
-                print(f"Truncating prompt to {self.max_prompt_length} tokens")
+                logger.info(f"Truncating prompt to {self.max_prompt_length} tokens")
             prompt_ids = prompt_ids[:, -self.max_prompt_length :]
             prompt_mask = prompt_mask[:, -self.max_prompt_length :]
 
         if self.max_completion_length is not None:
             if completion_ids.shape[1] > self.max_completion_length:
-                print(f"Truncating completion to {self.max_completion_length} tokens")
+                logger.info(
+                    f"Truncating completion to {self.max_completion_length} tokens"
+                )
             completion_ids = completion_ids[:, : self.max_completion_length]
             completion_mask = completion_mask[:, : self.max_completion_length]
 
@@ -225,7 +210,7 @@ class ArborGRPOTrainer(GRPOTrainer):
         prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1)
         attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)  # (B, P+C)
 
-        print(
+        logger.info(
             f"prompt_completion_ids.shape (after truncation, if enabled): {prompt_completion_ids.shape}, prompt_ids.shape: {prompt_ids.shape}, completion_ids.shape: {completion_ids.shape}"
         )
 
@@ -354,238 +339,14 @@ class ArborGRPOTrainer(GRPOTrainer):
 class LastStepTimeCallback(TrainerCallback):
     "A callback that prints a message at the beginning of training"
 
-    def on_step_end(self, args, state, control, **kwargs):
-        global last_step_time
-        print(f"Time since last step: {time_since_last_step()}")
-        last_step_time = time.time()
-
-
-class WeightUpdateCallback(TrainerCallback):
-    """A callback that sends weight update completion status after each step"""
-
-    def __init__(self):
-        self.comms_handler = None
-        self.trainer = None
-
-    def set_comms_handler(self, comms_handler: ArborScriptCommsHandler):
-        self.comms_handler = comms_handler
-
-    def set_trainer(self, trainer):
-        self.trainer = trainer
+    def __init__(self, ingestion_monitor: IngestionMonitor):
+        self.ingestion_monitor = ingestion_monitor
 
     def on_step_end(self, args, state, control, **kwargs):
-        if self.comms_handler and self.comms_handler.is_main_process and self.trainer:
-            if state.global_step != self.trainer._last_loaded_step:
-                print("Updating inference model...")
-                self.comms_handler.send_status({"status": "weight_update_start"})
-                self.trainer._move_model_to_vllm()
-                self.trainer._last_loaded_step = state.global_step
-                print("[DEBUG] Sending weight update completion status")
-                self.comms_handler.send_status({"status": "weight_update_complete"})
-
-
-class BlockingQueueDataset(Dataset):
-    def __init__(
-        self,
-        accelerator: Accelerator,
-        comms_handler: ArborScriptCommsHandler,
-        size=10_000,  # Just a random number
-        maxsize=100,
-    ):
-        self.size = size
-        self.accelerator = accelerator
-        self.comms_handler = comms_handler
-        self.get_cached_data = lru_cache(maxsize=maxsize)(self._get_data)
-        self.completion_counters = {}
-
-    def __len__(self):
-        return self.size
-
-    def _get_data(self, idx):
-        rank = self.accelerator.process_index
-        world_size = self.accelerator.num_processes
-
-        if self.accelerator.is_main_process:
-            global last_queue_pop_time
-            last_queue_pop_time = time.time()
-
-        if idx not in self.completion_counters:
-            self.completion_counters[idx] = 0
-
-        try:
-            new_data = self.comms_handler.receive_data()
-
-        except Exception as e:
-            print(f"[rank {rank}] Error receiving data: {e}")
-            new_data = None
-
-        return new_data
-
-    def __getitem__(self, idx):
-        data = self.get_cached_data(idx)
-        # Create hash of data to detect if processes are using the same idx for the same data
-        data_hash = format(abs(hash(str(data))) % (16**8), "08x")
-
-        if data is None:
-            return None
-
-        counter = self.completion_counters.get(idx, 0)
-        item = data[counter]
-        self.completion_counters[idx] = (counter + 1) % len(data)
-        return item
-
-
-class CommandMonitor:
-    def __init__(
-        self,
-        comms_handler: ArborScriptCommsHandler,
-        trainer: ArborGRPOTrainer,
-        base_model_name: str,
-    ):
-        self.comms_handler = comms_handler
-        self.trainer = trainer
-        self.base_model_name = base_model_name
-        self.command_thread = threading.Thread(
-            target=self._monitor_commands, daemon=True
+        logger.info(
+            f"Time since last step: {self.ingestion_monitor.time_since_last_step()}"
         )
-        self.command_thread.start()
-
-    def _monitor_commands(self):
-        """Background thread that monitors for commands from the server."""
-        if not self.comms_handler:
-            return
-        try:
-            for command in self.comms_handler.receive_command():
-                print(f"Main process received command: {command}")
-                if (
-                    command.get("command") == "save_model"
-                    and self.trainer.accelerator.is_main_process
-                ):
-                    print(
-                        f"[Training Script] Instructed to save model at {self.trainer.args.output_dir}"
-                    )
-                    while (
-                        time_since_last_step() <= 10
-                        or get_time_since_last_queue_pop() <= 10
-                    ):
-                        print(f"Waiting for steps to finish")
-                        print(
-                            f"Time since last step: {time_since_last_step():.1f} (needs to be >= 10)"
-                        )
-                        print(
-                            f"Time since last queue pop: {get_time_since_last_queue_pop():.1f} (needs to be >= 10)"
-                        )
-                        time.sleep(5)
-                    print("[Training Script] Saving model...")
-                    if self.trainer.peft_config:
-                        self.trainer.save_model(
-                            output_dir=self.trainer.args.output_dir + "/adapter/"
-                        )
-                        _model_to_merge = AutoPeftModelForCausalLM.from_pretrained(
-                            self.trainer.args.output_dir + "/adapter/",
-                            config=self.trainer.peft_config,
-                        )
-                        merged_model = _model_to_merge.merge_and_unload()
-                        merged_model.save_pretrained(
-                            self.trainer.args.output_dir,
-                            safe_serialization=True,
-                        )
-                        self.trainer.processing_class.save_pretrained(
-                            self.trainer.args.output_dir
-                        )
-                    else:
-                        self.trainer.save_model()
-
-                    print("[Training Script] Model saved")
-                    self.comms_handler.send_status(
-                        {
-                            "status": "model_saved",
-                            "output_dir": self.trainer.args.output_dir,
-                        }
-                    )
-                elif command.get("command") == "save_checkpoint":
-                    print(
-                        f"[Training Script] Instructed to save checkpoint {command.get('checkpoint_name')}"
-                    )
-                    while (
-                        time_since_last_step() <= 10
-                        or get_time_since_last_queue_pop() <= 10
-                    ):
-                        print(f"Waiting for steps to finish")
-                        print(
-                            f"Time since last step: {time_since_last_step():.1f} (needs to be >= 10)"
-                        )
-                        print(
-                            f"Time since last queue pop: {get_time_since_last_queue_pop():.1f} (needs to be >= 10)"
-                        )
-                        time.sleep(5)
-                    if self.trainer.peft_config:
-                        self.trainer.save_model(
-                            output_dir=self.trainer.args.output_dir
-                            + f"/checkpoints/{command.get('checkpoint_name')}/adapter/"
-                        )
-                        _model_to_merge = AutoPeftModelForCausalLM.from_pretrained(
-                            self.trainer.args.output_dir
-                            + f"/checkpoints/{command.get('checkpoint_name')}/adapter/",
-                            config=self.trainer.peft_config,
-                        )
-                        merged_model = _model_to_merge.merge_and_unload()
-                        merged_model.save_pretrained(
-                            self.trainer.args.output_dir
-                            + f"/checkpoints/{command.get('checkpoint_name')}/",
-                            safe_serialization=True,
-                        )
-                        self.trainer.processing_class.save_pretrained(
-                            self.trainer.args.output_dir
-                            + f"/checkpoints/{command.get('checkpoint_name')}/"
-                        )
-                    else:
-                        self.trainer.save_model(
-                            output_dir=self.trainer.args.output_dir
-                            + f"/checkpoints/{command.get('checkpoint_name')}/"
-                        )
-
-                    # Copy checkpoint files to root output directory
-                    checkpoint_dir = (
-                        self.trainer.args.output_dir
-                        + f"/checkpoints/{command.get('checkpoint_name')}/"
-                    )
-                    root_dir = self.trainer.args.output_dir
-
-                    # Copy all files from checkpoint dir to root dir, overwriting if they exist
-                    # (effectively saves the checkpoint to the output directory)
-                    for item in os.listdir(checkpoint_dir):
-                        src = os.path.join(checkpoint_dir, item)
-                        dst = os.path.join(root_dir, item)
-                        if os.path.isdir(src):
-                            if os.path.exists(dst):
-                                shutil.rmtree(dst)
-                            shutil.copytree(src, dst)
-                        else:
-                            shutil.copy2(src, dst)
-
-                    self.comms_handler.send_status(
-                        {
-                            "status": "checkpoint_saved",
-                            "checkpoint_name": command.get("checkpoint_name"),
-                            "output_dir": self.trainer.args.output_dir
-                            + f"/checkpoints/{command.get('checkpoint_name')}/",
-                        }
-                    )
-                    self.comms_handler.send_status(
-                        {
-                            "status": "model_saved",
-                            "output_dir": self.trainer.args.output_dir,
-                        }
-                    )
-                elif command.get("command") == "terminate":
-                    print("TERMINATED")
-                    self.trainer.accelerator.end_training()
-                    self.comms_handler.send_status({"status": "terminated"})
-
-        except Exception as e:
-            print(e)
-            self.comms_handler.send_status({"status": "error", "error": str(e)})
+        self.ingestion_monitor.set_last_step_time()
 
 
 def main():
@@ -679,7 +440,7 @@ def main():
             # Need to set subscription for PUB/SUB pattern
             server_comms_handler.status_socket.setsockopt_string(zmq.SUBSCRIBE, "")
             for status in server_comms_handler.receive_status():
-                print(f"Status: {status}")
+                logger.info(f"Status: {status}")
 
         status_listener_thread = threading.Thread(target=status_listener, daemon=True)
         status_listener_thread.start()
@@ -693,7 +454,7 @@ def main():
         if "gradient_checkpointing_kwargs" in trl_train_args and arbor_train_args.get(
             "lora", False
         ):
-            print(
+            logger.info(
                 "Setting gradient_checkpointing_kwargs to use_reentrant=False for LORA training"
             )
             trl_train_args["gradient_checkpointing_kwargs"] = {
@@ -703,7 +464,7 @@ def main():
 
         lora_config = None
         if arbor_train_args.get("lora", False):
-            print("Using LORA for PEFT")
+            logger.info("Using LORA for PEFT")
             lora_config = LoraConfig(
                 r=16,
                 lora_alpha=64,
@@ -721,6 +482,12 @@ def main():
                 inference_mode=False,
             )
 
+        if "report_to" in trl_train_args and trl_train_args["report_to"] == "wandb":
+            import wandb
+
+            if "wandb_kwargs" in arbor_train_args and arbor_train_args["wandb_kwargs"]:
+                wandb.init(**arbor_train_args["wandb_kwargs"])
+
         training_args = GRPOConfig(
             dataloader_num_workers=0,
             shuffle_dataset=False,
@@ -728,15 +495,24 @@ def main():
             **trl_train_args,
         )
 
-        weight_update_callback = WeightUpdateCallback()
+        # Create ingestion monitor
+        ingestion_monitor = IngestionMonitor()
+
+        train_dataset = BlockingRotatingQueueDataset(
+            ingestion_monitor=ingestion_monitor,
+        )
+
+        weight_update_callback = WeightUpdateCallback(
+            ingestion_monitor=ingestion_monitor,
+        )
+
         trainer = ArborGRPOTrainer(
             model=args.model,
             args=training_args,
-            train_dataset=BlockingQueueDataset(None, None),
-            callbacks=[LastStepTimeCallback(), weight_update_callback],
+            train_dataset=train_dataset,
+            callbacks=[LastStepTimeCallback(ingestion_monitor), weight_update_callback],
             peft_config=lora_config,
             vllm_group_port=args.vllm_group_port,
-            **arbor_train_args,
         )
         # Create client handler
         comms_handler = ArborScriptCommsHandler(
@@ -748,48 +524,53 @@ def main():
             handshake_port=args.handshake_port,
             is_main_process=trainer.accelerator.is_main_process,
         )
+
+        train_dataset.set_comms_handler(comms_handler)
+        train_dataset.set_accelerator(trainer.accelerator)
+
         weight_update_callback.set_comms_handler(comms_handler)
         weight_update_callback.set_trainer(trainer)
         trainer.comms_handler = comms_handler
 
-        # Initialize the dataset with the actual accelerator
-        trainer.train_dataset = BlockingQueueDataset(
-            accelerator=trainer.accelerator,
-            comms_handler=trainer.comms_handler,
-        )
-
         command_monitor = CommandMonitor(
             comms_handler=comms_handler,
             trainer=trainer,
             base_model_name=args.model,
+            ingestion_monitor=ingestion_monitor,
         )
+        command_monitor.start()
 
         # Add signal handlers for graceful shutdown
         def signal_handler(signum, frame):
-            print(f"\nReceived signal {signum}. Initiating graceful shutdown...")
-            print("Ending training...")
+            logger.info(f"\nReceived signal {signum}. Initiating graceful shutdown...")
+            logger.info("Ending training...")
             trainer.accelerator.end_training()
-            print("Closing communications...")
+            logger.info("Closing communications...")
             comms_handler.close()
             sys.exit(0)
 
         signal.signal(signal.SIGINT, signal_handler)
         signal.signal(signal.SIGTERM, signal_handler)
 
-        print("Training...")
-        trainer.train()
+        logger.info("Starting training...")
+        try:
+            trainer.train()
+        except Exception as e:
+            logger.error(f"Error during training: {e}")
+            logger.error(f"Error type: {type(e).__name__}")
+            raise
 
     except KeyboardInterrupt:
-        print("\nReceived interrupt, shutting down...")
+        logger.info("\nReceived interrupt, shutting down...")
     except Exception as e:
-        print(f"Error: {e}")
+        logger.error(f"Error: {e}")
         comms_handler.send_status({"status": "error", "error": str(e)})
         raise e
     finally:
-        print("Cleaning up resources...")
+        logger.info("Cleaning up resources...")
         trainer.accelerator.end_training()
         comms_handler.close()
-        print("Cleanup complete")
+        logger.info("Cleanup complete")
 
 
 if __name__ == "__main__":