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

arbor/server/services/scripts/utils/comms_monitors.py

@@ -0,0 +1,169 @@
+import os
+import shutil
+import threading
+import time
+from typing import Callable, Optional
+
+from peft import AutoPeftModelForCausalLM
+from transformers import Trainer
+
+from arbor.server.services.comms.comms import ArborScriptCommsHandler
+
+
+class CommandMonitor:
+    def __init__(
+        self,
+        comms_handler: ArborScriptCommsHandler,
+        trainer: Trainer,
+        base_model_name: str,
+        ingestion_monitor: Optional["IngestionMonitor"] = None,
+    ):
+        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
+        )
+        self.ingestion_monitor = ingestion_monitor
+
+    def start(self):
+        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 self.ingestion_monitor and (
+                        self.ingestion_monitor.time_since_last_step() <= 10
+                        or self.ingestion_monitor.time_since_last_queue_pop() <= 10
+                    ):
+                        print(f"Waiting for steps to finish")
+                        if self.ingestion_monitor:
+                            print(
+                                f"Time since last step: {self.ingestion_monitor.time_since_last_step():.1f} (needs to be >= 10)"
+                            )
+                            print(
+                                f"Time since last queue pop: {self.ingestion_monitor.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 self.ingestion_monitor and (
+                        self.ingestion_monitor.time_since_last_step() <= 10
+                        or self.ingestion_monitor.time_since_last_queue_pop() <= 10
+                    ):
+                        print(f"Waiting for steps to finish")
+                        if self.ingestion_monitor:
+                            print(
+                                f"Time since last step: {self.ingestion_monitor.time_since_last_step():.1f} (needs to be >= 10)"
+                            )
+                            print(
+                                f"Time since last queue pop: {self.ingestion_monitor.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)})

arbor/server/services/scripts/utils/dataset.py

@@ -0,0 +1,176 @@
+import json
+import logging
+import time
+from functools import lru_cache
+from typing import Any, Callable, Dict, List, Optional
+
+from accelerate import Accelerator
+from datasets import Dataset as HuggingFaceDataset
+from torch.utils.data import Dataset as TorchDataset
+
+from arbor.server.services.comms.comms import ArborScriptCommsHandler
+
+logger = logging.getLogger(__name__)
+
+
+class BlockingRotatingQueueDataset(TorchDataset):
+    def __init__(
+        self,
+        size=10_000,  # Just a random number
+        maxsize=100,
+        ingestion_monitor: Optional["IngestionMonitor"] = None,
+    ):
+        self.size = size
+        # Use a regular cache dict instead of lru_cache to avoid unhashable type issues
+        self._data_cache = {}
+        self._cache_maxsize = maxsize
+        self.completion_counters = {}
+        self.ingestion_monitor = ingestion_monitor
+        self.accelerator = None
+        self.comms_handler = None
+
+    def set_accelerator(self, accelerator: Accelerator):
+        self.accelerator = accelerator
+
+    def set_comms_handler(self, comms_handler: ArborScriptCommsHandler):
+        self.comms_handler = comms_handler
+
+    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 and self.ingestion_monitor:
+            self.ingestion_monitor.set_last_queue_pop_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}")
+            if "unhashable" in str(e):
+                print(
+                    f"[rank {rank}] DEBUGGING: Unhashable type error in data reception"
+                )
+                print(
+                    f"[rank {rank}] This might be related to caching or data structure issues"
+                )
+            new_data = None
+
+        return new_data
+
+    def get_cached_data(self, idx):
+        """Get data with simple dictionary caching instead of lru_cache"""
+        if idx in self._data_cache:
+            return self._data_cache[idx]
+
+        # If cache is full, clear oldest entries (simple FIFO)
+        if len(self._data_cache) >= self._cache_maxsize:
+            # Remove first half of cache entries
+            keys_to_remove = list(self._data_cache.keys())[: self._cache_maxsize // 2]
+            for key in keys_to_remove:
+                del self._data_cache[key]
+
+        # Get new data and cache it
+        data = self._get_data(idx)
+        self._data_cache[idx] = data
+        return data
+
+    def __getitem__(self, idx):
+        print(f"Getting item {idx}")
+        data = self.get_cached_data(idx)
+
+        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 BlockingQueueDataset(HuggingFaceDataset):
+    def __init__(
+        self,
+        ingestion_monitor: Optional["IngestionMonitor"] = None,
+    ):
+        self._buffer: List[Dict[str, Any]] = []
+        self._logger = logging.getLogger(__name__)
+        self.ingestion_monitor = ingestion_monitor
+
+    def set_accelerator(self, accelerator: Accelerator):
+        self.accelerator = accelerator
+
+    def set_comms_handler(self, comms_handler: ArborScriptCommsHandler):
+        self.comms_handler = comms_handler
+
+    def __len__(self) -> int:
+        return 1_000_000
+
+    def _fill_buffer(self, target_size: int) -> None:
+        while len(self._buffer) < target_size:
+            try:
+                if self.comms_handler is None:
+                    raise ValueError("comms_handler is not initialized")
+
+                group = self.comms_handler.receive_data()
+
+                if group is not None:
+                    self._logger.debug("Received group from comms handler")
+                    for trajectory in group:
+                        trajectory_copy = json.loads(json.dumps(trajectory))
+                        for item in trajectory:
+                            item["trajectory"] = trajectory_copy
+                            self._buffer.append(item)
+
+            except Exception as e:
+                if "Context was terminated" in str(e):
+                    self._logger.error(
+                        "ZMQ context was terminated while filling buffer"
+                    )
+                    raise RuntimeError("ZMQ context was terminated") from e
+                self._logger.warning(f"Error receiving data: {e}")
+                continue
+
+    def _transform_batch(self, items: List[Dict[str, Any]]) -> Dict[str, List[Any]]:
+        if not items:
+            raise ValueError("Cannot transform empty batch")
+
+        return {key: [item[key] for item in items] for key in items[0].keys()}
+
+    def __getitem__(self, idx: List[int]) -> Dict[str, List[Any]]:
+        if self.accelerator is None:
+            self._logger.error("Accelerator not initialized")
+            raise ValueError("Accelerator must be initialized before getting items")
+        if self.comms_handler is None:
+            self._logger.error("Comms handler not initialized")
+            raise ValueError("Comms handler must be initialized before getting items")
+
+        batch_size = len(idx)
+        if batch_size == 0:
+            raise ValueError("Batch size must be greater than 0")
+
+        try:
+            self._fill_buffer(batch_size)
+
+            if len(self._buffer) < batch_size:
+                raise RuntimeError(
+                    f"Not enough items in buffer (got {len(self._buffer)}, need {batch_size})"
+                )
+
+            batch_items = self._buffer[:batch_size]
+            self._buffer = self._buffer[batch_size:]
+
+            if self.ingestion_monitor:
+                self.ingestion_monitor.set_last_queue_pop_time()
+
+            return self._transform_batch(batch_items)
+
+        except Exception as e:
+            self._logger.error(f"Error getting batch: {e}")
+            raise

arbor/server/services/scripts/utils/ingestion_monitor.py

@@ -0,0 +1,35 @@
+import time
+from typing import Optional
+
+
+class IngestionMonitor:
+    """Monitor for tracking timing of training steps and data ingestion (queue pops)"""
+
+    def __init__(self):
+        self._last_step_time: Optional[float] = None
+        self._last_queue_pop_time: Optional[float] = None
+
+    def time_since_last_step(self) -> float:
+        """Get time elapsed since last training step"""
+        if self._last_step_time is None:
+            return float("inf")
+        return time.time() - self._last_step_time
+
+    def time_since_last_queue_pop(self) -> float:
+        """Get time elapsed since last queue pop"""
+        if self._last_queue_pop_time is None:
+            return float("inf")
+        return time.time() - self._last_queue_pop_time
+
+    def set_last_queue_pop_time(self, timestamp: Optional[float] = None) -> None:
+        """Set the last queue pop time"""
+        self._last_queue_pop_time = timestamp if timestamp is not None else time.time()
+
+    def set_last_step_time(self, timestamp: Optional[float] = None) -> None:
+        """Set the last step time"""
+        self._last_step_time = timestamp if timestamp is not None else time.time()
+
+    def reset(self) -> None:
+        """Reset all timing data"""
+        self._last_step_time = None
+        self._last_queue_pop_time = None

arbor/server/services/scripts/utils/mock_server.py

@@ -0,0 +1,124 @@
+## Mock arbor sending over data for testing
+import threading
+import time
+
+import zmq
+
+from arbor.server.services.comms.comms import ArborServerCommsHandler
+
+group_example = [  # Entire group of trajectories
+    [  # Trajectory with different modules
+        {  # geography module
+            "messages": [{"role": "user", "content": "What is the capital of France?"}],
+            "completion": [{"role": "assistant", "content": "Paris"}],
+            "advantage": 0.9,
+        },
+        {  # math module
+            "messages": [{"role": "user", "content": "What is 2 * 2 + 2?"}],
+            "completion": [{"role": "assistant", "content": "6"}],
+            "advantage": 0.8,
+        },
+        {  # car module
+            "messages": [
+                {"role": "user", "content": "When did the first honda civic come out?"}
+            ],
+            "completion": [{"role": "assistant", "content": "1973"}],
+            "advantage": 0.7,
+        },
+    ],
+    [  # Trajectory with different modules
+        {  # geography module
+            "messages": [
+                {"role": "user", "content": "What is the capital of Germany?"}
+            ],
+            "completion": {
+                "role": "assistant",
+                "content": "Berlin is the capital of Germany",
+            },
+            "advantage": 0.1,
+        },
+        {  # math module
+            "messages": [{"role": "user", "content": "What is 2 + 2?"}],
+            "completion": [{"role": "assistant", "content": "3"}],
+            "advantage": 0.2,
+        },
+    ],
+]
+
+
+def flatten_batch(batch):
+    return [item for sublist in batch for item in sublist]
+
+
+def debug_data_generator(server_comms_handler):
+    idx = 0
+    while True:
+        print(f"Sending group:")  # Debug print
+        server_comms_handler.send_data(group_example)
+        idx += 1
+        time.sleep(1)
+
+        if idx >= 100:
+            server_comms_handler.send_command({"command": "save_model"})
+
+
+def status_listener(server_comms_handler):
+    # 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}")
+
+
+if __name__ == "__main__":
+    server_comms_handler = ArborServerCommsHandler(
+        host="localhost",
+    )
+
+    # Get available ports from the server comms handler
+    command_port = server_comms_handler.command_port
+    status_port = server_comms_handler.status_port
+    data_port = server_comms_handler.data_port
+    broadcast_port = server_comms_handler.broadcast_port
+    handshake_port = server_comms_handler.handshake_port
+
+    # Print the command that would be used to connect to this mock server
+    print("\nTo connect to this mock server, run the following command:")
+    print(
+        f"CUDA_VISIBLE_DEVICES=2 python arbor/server/services/scripts/mmgrpo_training.py \\"
+    )
+    print(f"    --debug \\")
+    print(f"    --command_port {command_port} \\")
+    print(f"    --status_port {status_port} \\")
+    print(f"    --data_port {data_port} \\")
+    print(f"    --broadcast_port {broadcast_port} \\")
+    print(f"    --handshake_port {handshake_port} \\")
+    print(f"    --vllm_group_port 0 \\")
+    print(f"    --vllm_port 0 \\")
+    print(f"    --model Qwen/Qwen3-0.6B \\")
+    print(f'    --trl_train_kwargs \'{{"output_dir": ".", "report_to": "none"}}\'')
+    print(
+        "\nThis mock server will simulate sending training data to the training process."
+    )
+    print("Press Ctrl+C to exit the mock server.\n")
+
+    server_comms_handler.wait_for_clients(1)
+
+    debug_thread = threading.Thread(
+        target=debug_data_generator, args=(server_comms_handler,), daemon=True
+    )
+    debug_thread.start()
+
+    status_listener_thread = threading.Thread(
+        target=status_listener, args=(server_comms_handler,), daemon=True
+    )
+    status_listener_thread.start()
+
+    try:
+        print("Mock server started and waiting for training process to connect...")
+        while True:
+            time.sleep(1)
+    except KeyboardInterrupt:
+        print("\nShutting down mock server...")
+    finally:
+        server_comms_handler.close()
+        print("Mock server shutdown complete.")

arbor/server/services/training_manager.py

@@ -6,14 +6,14 @@ from datetime import datetime
 from pathlib import Path
 
 from arbor.server.api.models.schemas import FineTuneRequest
-from arbor.server.core.config import Settings
+from arbor.server.core.config import Config
 from arbor.server.services.file_manager import FileManager
 from arbor.server.services.job_manager import Job, JobEvent, JobStatus
 
 
 class TrainingManager:
-    def __init__(self, settings: Settings):
-        self.settings = settings
+    def __init__(self, config: Config):
+        self.config = config
 
     def make_output_dir(self, request: FineTuneRequest):
         model_name = request.model.split("/")[-1].lower()
@@ -24,7 +24,7 @@ class TrainingManager:
         )
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
         name = f"ft:{model_name}:{suffix}:{timestamp}"
-        return name, str(Path(self.settings.STORAGE_PATH).resolve() / "models" / name)
+        return name, str(Path(self.config.STORAGE_PATH).resolve() / "models" / name)
 
     def find_train_args_sft(self, request: FineTuneRequest, file_manager: FileManager):
         file = file_manager.get_file(request.training_file)