arbor-ai 0.1.12__py3-none-any.whl → 0.1.14__py3-none-any.whl

arbor/server/api/models/schemas.py

@@ -199,10 +199,16 @@ class GRPOConfigRequest(BaseModel):
     bf16: Optional[bool] = None
     scale_rewards: Optional[bool] = None
     max_grad_norm: Optional[float] = None
+    report_to: Optional[str] = None
+    log_completions: Optional[bool] = None
+    logging_steps: Optional[int] = None
+    mask_truncated_completions: Optional[bool] = None
+    # Arbor specific
+    max_context_length: Optional[int] = None
     lora: Optional[bool] = None
-    update_interval: Optional[int] = None
     # To name the run
     suffix: Optional[str] = None
+    generation_batch_size: Optional[int] = None
 
 
 class GRPOConfigResponse(BaseModel):
@@ -216,8 +222,23 @@ class GRPOTerminateRequest(BaseModel):
 class GRPOTerminateResponse(BaseModel):
     status: str
     current_model: str
+    checkpoints: Optional[dict[str, str]] = None
+    last_checkpoint: Optional[str] = None
 
 
 class GRPOStepResponse(BaseModel):
     status: str
     current_model: str
+    checkpoints: dict[str, str]
+    last_checkpoint: Optional[str] = None
+
+
+class GRPOCheckpointRequest(BaseModel):
+    checkpoint_name: str
+
+
+class GRPOCheckpointResponse(BaseModel):
+    status: str
+    current_model: str
+    checkpoints: dict[str, str]
+    last_checkpoint: str

arbor/server/api/routes/grpo.py

@@ -4,6 +4,8 @@ import subprocess
 from fastapi import APIRouter, BackgroundTasks, Request
 
 from arbor.server.api.models.schemas import (
+    GRPOCheckpointRequest,
+    GRPOCheckpointResponse,
     GRPOConfigRequest,
     GRPOConfigResponse,
     GRPORequest,
@@ -31,17 +33,24 @@ def run_grpo_step(
     inference_manager = request.app.state.inference_manager
     grpo_manager = request.app.state.grpo_manager
 
-    current_model = grpo_manager.grpo_step(grpo_request, inference_manager)
+    step_data = grpo_manager.grpo_step(grpo_request, inference_manager)
 
-    return GRPOStepResponse(status="success", current_model=current_model)
+    return GRPOStepResponse(status="success", **step_data)
 
 
 @router.post("/update_model", response_model=GRPOStepResponse)
 def update_model(request: Request):
     grpo_manager = request.app.state.grpo_manager
     inference_manager = request.app.state.inference_manager
-    current_model = grpo_manager.update_model(request, inference_manager)
-    return GRPOStepResponse(status="success", current_model=current_model)
+    update_model_data = grpo_manager.update_model(request, inference_manager)
+    return GRPOStepResponse(status="success", **update_model_data)
+
+
+@router.post("/checkpoint", response_model=GRPOCheckpointResponse)
+def checkpoint(request: Request, grpo_checkpoint_request: GRPOCheckpointRequest):
+    grpo_manager = request.app.state.grpo_manager
+    checkpoint_data = grpo_manager.checkpoint(grpo_checkpoint_request)
+    return GRPOCheckpointResponse(status="success", **checkpoint_data)
 
 
 @router.post("/terminate", response_model=GRPOTerminateResponse)
@@ -50,5 +59,5 @@ def terminate_grpo(request: Request):
     grpo_manager = request.app.state.grpo_manager
     inference_manager = request.app.state.inference_manager
 
-    final_model = grpo_manager.terminate(inference_manager)
-    return GRPOTerminateResponse(status="success", current_model=final_model)
+    terminate_data = grpo_manager.terminate(inference_manager)
+    return GRPOTerminateResponse(status="success", **terminate_data)

arbor/server/services/grpo_manager.py

@@ -13,7 +13,11 @@ from datetime import datetime
 from pathlib import Path
 from typing import Optional
 
-from arbor.server.api.models.schemas import GRPOConfigRequest, GRPORequest
+from arbor.server.api.models.schemas import (
+    GRPOCheckpointRequest,
+    GRPOConfigRequest,
+    GRPORequest,
+)
 from arbor.server.core.config import Settings
 from arbor.server.services.comms.comms import ArborServerCommsHandler
 from arbor.server.services.inference_manager import InferenceManager
@@ -28,7 +32,10 @@ class GRPOManager:
         self.server_comms_handler = None
         self.status_thread = None
         self.model_saved_and_reload_requested = False
+        self.saving_checkpoint = False
 
+        self.checkpoints = {}
+        self.last_checkpoint = None
         self.data_count = 0
         self.last_inference_update = 0
         # Set up signal handler
@@ -86,12 +93,17 @@ class GRPOManager:
             "bf16",
             "scale_rewards",
             "max_grad_norm",
+            "report_to",
+            "log_completions",
+            "logging_steps",
+            "generation_batch_size",
+            "mask_truncated_completions",
         ]
         trl_train_kwargs = {
             key: train_kwargs[key] for key in trl_keys if key in train_kwargs
         }
 
-        arbor_keys = ["update_interval", "lora"]
+        arbor_keys = ["max_context_length", "lora"]
         arbor_train_kwargs = {
             key: train_kwargs[key] for key in arbor_keys if key in train_kwargs
         }
@@ -119,6 +131,8 @@ class GRPOManager:
         # Start the training process with ZMQ ports
         my_env = os.environ.copy()
         my_env["CUDA_VISIBLE_DEVICES"] = self.settings.arbor_config.training.gpu_ids
+        # WandB can block the training process for login, so we silence it
+        my_env["WANDB_SILENT"] = "true"
 
         num_processes = self.settings.arbor_config.training.gpu_ids.count(",") + 1
 
@@ -209,6 +223,12 @@ class GRPOManager:
 
         # Launch the inference server
         print("Launching inference server...")
+        # launch_kwargs = {
+        #     k: v for k, v in arbor_train_kwargs.items() if k in ["max_context_length"]
+        # }
+        inference_manager.launch_kwargs["max_context_length"] = arbor_train_kwargs.get(
+            "max_context_length", None
+        )
         inference_manager.launch(self.current_model)
 
     def _handle_status_updates(self, inference_manager: InferenceManager):
@@ -228,6 +248,12 @@ class GRPOManager:
                         self.model_saved_and_reload_requested = False
                         self.current_model = status["output_dir"]
                         print("Model update complete")
+                elif status["status"] == "checkpoint_saved":
+                    print("Received checkpoint saved status")
+                    self.checkpoints[status["checkpoint_name"]] = status["output_dir"]
+                    self.last_checkpoint = status["checkpoint_name"]
+                    self.saving_checkpoint = False
+                    print("Checkpoint saved")
                 elif status["status"] == "error":
                     print(f"Training error: {status.get('error', 'Unknown error')}")
                 elif status["status"] == "terminated":
@@ -249,6 +275,10 @@ class GRPOManager:
             )
             time.sleep(5)
 
+        while self.saving_checkpoint:
+            print("Saving checkpoint, pausing GRPO steps until checkpoint is saved...")
+            time.sleep(5)
+
         try:
             # Send the batch to the training process
             self.server_comms_handler.send_data(request.batch)
@@ -256,12 +286,11 @@ class GRPOManager:
         except Exception as e:
             print(f"Failed to send batch to training process: {e}")
 
-        # We tell the script to save the model. The script will let us know when it's done via the status update handler
-        # Then we'll actually run the update_model function in the inference manager and finally update the last_inference_update variable
-        # if self._should_update_model():
-        #     self.server_comms_handler.send_command({"command": "save_model"})
-
-        return self.current_model
+        return {
+            "current_model": self.current_model,
+            "checkpoints": self.checkpoints,
+            "last_checkpoint": self.last_checkpoint,
+        }
 
     def update_model(self, request, inference_manager: InferenceManager):
         if inference_manager._session:
@@ -286,18 +315,41 @@ class GRPOManager:
                 "Waiting for model to be saved and reloaded... This usually takes 20-30 seconds"
             )
             time.sleep(5)
-        return self.current_model
+        return {
+            "current_model": self.current_model,
+            "checkpoints": self.checkpoints,
+            "last_checkpoint": self.last_checkpoint,
+        }
+
+    def checkpoint(self, request: GRPOCheckpointRequest):
+        self.saving_checkpoint = True
+        self.server_comms_handler.send_command(
+            {"command": "save_checkpoint", "checkpoint_name": request.checkpoint_name}
+        )
+        while self.saving_checkpoint:
+            print("Waiting for checkpoint to be saved...")
+            time.sleep(5)
+        return {
+            "current_model": self.current_model,
+            "checkpoints": self.checkpoints,
+            "last_checkpoint": self.last_checkpoint,
+        }
 
     def terminate(self, inference_manager: InferenceManager):
         """Clean up resources and save the final model."""
+        termination_data = {
+            "current_model": self.current_model,
+            "checkpoints": self.checkpoints,
+            "last_checkpoint": self.last_checkpoint,
+        }
         try:
             # Stop the inference server
             if inference_manager.process is not None:
                 inference_manager.kill()
 
             # Send termination command through REQ socket
-            # self.server_comms_handler.send_broadcast({"message": "terminate"})
-            self.training_process.terminate()
+            self.server_comms_handler.send_broadcast({"message": "terminate"})
+            # self.training_process.terminate()
             print("Waiting for training process to finish")
 
             # Wait for training process to finish
@@ -336,17 +388,13 @@ class GRPOManager:
                     )
                 output_dir = self.train_kwargs["output_dir"]
                 self.train_kwargs = None
-                return output_dir
             else:
                 print("Training terminated, no output directory specified")
                 self.train_kwargs = None
-                return None
+
+        return termination_data
 
     def _should_update_model(self):
-        # return (
-        #     self.data_count - self.last_inference_update
-        #     >= self.train_kwargs["update_interval"]
-        # )
         return self.model_saved_and_reload_requested
 
 

arbor/server/services/inference/sgl_router_launch_server.py

@@ -0,0 +1,226 @@
+import argparse
+import copy
+import json
+import logging
+import multiprocessing as mp
+import os
+import random
+import signal
+import sys
+import time
+from typing import List
+
+import requests
+import zmq
+from setproctitle import setproctitle
+from sglang.srt.entrypoints.http_server import launch_server
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils import is_port_available
+from sglang_router.launch_router import RouterArgs, launch_router
+
+
+def setup_logger():
+    logger = logging.getLogger("router")
+    logger.setLevel(logging.INFO)
+
+    formatter = logging.Formatter(
+        "[Router (Python)] %(asctime)s - %(levelname)s - %(message)s - %(filename)s:%(lineno)d",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+
+    handler = logging.StreamHandler()
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+
+    return logger
+
+
+logger = setup_logger()
+
+
+# Create new process group
+def run_server(server_args, dp_rank):
+    """
+    Note:
+
+    1. Without os.setpgrp(), all processes share the same PGID. When you press Ctrl+C, the terminal sends SIGINT to all processes in the group simultaneously.
+    This can cause leaf processes to terminate first, which messes up the cleaning order and produces orphaned processes.
+
+    Terminal (PGID=100)
+    └── Main Python Process (PGID=100)
+        └── Server Process 1 (PGID=100)
+            └── Scheduler 1
+            └── Detokenizer 1
+        └── Server Process 2 (PGID=100)
+            └── Scheduler 2
+            └── Detokenizer 2
+
+    2. With os.setpgrp(), the main Python process and its children are in a separate group. Now:
+
+    Terminal (PGID=100)
+    └── Main Python Process (PGID=200)
+        └── Server Process 1 (PGID=300)
+            └── Scheduler 1
+            └── Detokenizer 1
+        └── Server Process 2 (PGID=400)
+            └── Scheduler 2
+            └── Detokenizer 2
+    """
+    # create new process group
+    os.setpgrp()
+
+    setproctitle("sglang::server")
+    # Set SGLANG_DP_RANK environment variable
+    os.environ["SGLANG_DP_RANK"] = str(dp_rank)
+
+    launch_server(server_args)
+
+
+def launch_server_process(
+    server_args: ServerArgs, worker_port: int, dp_id: int
+) -> mp.Process:
+    """Launch a single server process with the given args and port."""
+    server_args = copy.deepcopy(server_args)
+    server_args.port = worker_port
+    server_args.base_gpu_id = dp_id * server_args.tp_size
+    server_args.dp_size = 1
+
+    proc = mp.Process(target=run_server, args=(server_args, dp_id))
+    proc.start()
+    return proc
+
+
+def wait_for_server_health(host: str, port: int, timeout: int = 300) -> bool:
+    """Wait for server to be healthy by checking /health endpoint."""
+    start_time = time.time()
+    url = f"http://{host}:{port}/health"
+
+    while time.time() - start_time < timeout:
+        try:
+            response = requests.get(url, timeout=5)
+            if response.status_code == 200:
+                return True
+        except requests.exceptions.RequestException:
+            pass
+        time.sleep(1)
+    return False
+
+
+def find_available_ports(base_port: int, count: int) -> List[int]:
+    """Find consecutive available ports starting from base_port."""
+    available_ports = []
+    current_port = base_port
+
+    while len(available_ports) < count:
+        if is_port_available(current_port):
+            available_ports.append(current_port)
+        current_port += random.randint(100, 1000)
+
+    return available_ports
+
+
+def cleanup_processes(processes: List[mp.Process]):
+    for process in processes:
+        logger.info(f"Terminating process group {process.pid}")
+        try:
+            os.killpg(process.pid, signal.SIGTERM)
+        except ProcessLookupError:
+            # Process group may already be terminated
+            pass
+
+    # Wait for processes to terminate
+    for process in processes:
+        process.join(timeout=5)
+        if process.is_alive():
+            logger.warning(
+                f"Process {process.pid} did not terminate gracefully, forcing kill"
+            )
+            try:
+                os.killpg(process.pid, signal.SIGKILL)
+            except ProcessLookupError:
+                pass
+
+    logger.info("All process groups terminated")
+
+
+def main():
+    # CUDA runtime isn't fork-safe, which can lead to subtle bugs or crashes
+    mp.set_start_method("spawn")
+
+    parser = argparse.ArgumentParser(
+        description="Launch SGLang router and server processes"
+    )
+
+    ServerArgs.add_cli_args(parser)
+    RouterArgs.add_cli_args(parser, use_router_prefix=True, exclude_host_port=True)
+    parser.add_argument(
+        "--router-dp-worker-base-port",
+        type=int,
+        default=31000,
+        help="Base port number for data parallel workers",
+    )
+    parser.add_argument(
+        "--worker-urls-port",
+        type=int,
+        help="Port number for worker URLs publisher",
+    )
+
+    args = parser.parse_args()
+    server_args = ServerArgs.from_cli_args(args)
+    router_args = RouterArgs.from_cli_args(args, use_router_prefix=True)
+
+    # Find available ports for workers
+    worker_ports = find_available_ports(
+        args.router_dp_worker_base_port, server_args.dp_size
+    )
+
+    # Start server processes
+    server_processes = []
+
+    for i, worker_port in enumerate(worker_ports):
+        logger.info(f"Launching DP server process {i} on port {worker_port}")
+        proc = launch_server_process(server_args, worker_port, i)
+        server_processes.append(proc)
+
+    signal.signal(signal.SIGINT, lambda sig, frame: cleanup_processes(server_processes))
+    signal.signal(
+        signal.SIGTERM, lambda sig, frame: cleanup_processes(server_processes)
+    )
+    signal.signal(
+        signal.SIGQUIT, lambda sig, frame: cleanup_processes(server_processes)
+    )
+
+    # Update router args with worker URLs
+    worker_urls = [f"http://{server_args.host}:{port}" for port in worker_ports]
+    router_args.worker_urls = worker_urls
+
+    # Publish worker URLs via ZMQ if port is specified
+    if args.worker_urls_port:
+        try:
+            context = zmq.Context()
+            socket = context.socket(zmq.PUB)
+            socket.bind(f"tcp://*:{args.worker_urls_port}")
+            # Give subscribers time to connect
+            time.sleep(0.1)
+            socket.send_json({"type": "worker_urls", "urls": worker_urls})
+            logger.info(
+                f"Published worker URLs via ZMQ on port {args.worker_urls_port}"
+            )
+            socket.close()
+            context.term()
+        except Exception as e:
+            logger.error(f"Failed to publish worker URLs via ZMQ: {e}")
+            cleanup_processes(server_processes)
+            sys.exit(1)
+
+    # Start the router
+    try:
+        launch_router(router_args)
+    except Exception as e:
+        logger.error(f"Failed to start router: {e}")
+        cleanup_processes(server_processes)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

arbor/server/services/inference_manager.py

@@ -1,6 +1,7 @@
 import asyncio
 import json
 import os
+import random
 import signal
 import socket
 import subprocess
@@ -12,6 +13,7 @@ from typing import Any, Dict, Optional
 
 import aiohttp
 import requests
+import zmq
 
 from arbor.server.core.config import Settings
 
@@ -27,6 +29,7 @@ class InferenceManager:
         self.current_model = None
         self.inference_count = 0
         self._session = None
+        self.worker_urls = []
         # Set up signal handler for graceful shutdown
         signal.signal(signal.SIGINT, self._signal_handler)
         signal.signal(signal.SIGTERM, self._signal_handler)
@@ -47,7 +50,12 @@ class InferenceManager:
     def is_server_restarting(self):
         return self.restarting
 
-    def launch(self, model: str, launch_kwargs: Optional[Dict[str, Any]] = None):
+    def launch(
+        self,
+        model: str,
+        launch_kwargs: Optional[Dict[str, Any]] = None,
+        max_retries: int = 3,
+    ):
         if self.is_server_running():
             print("Server is already launched.")
             return
@@ -59,77 +67,122 @@ class InferenceManager:
             if model.startswith(prefix):
                 model = model[len(prefix) :]
 
-        print(f"Grabbing a free port to launch an SGLang server for model {model}")
-        port = get_free_port()
-        timeout = launch_kwargs.get("timeout", 1800)
-        my_env = os.environ.copy()
-        my_env["CUDA_VISIBLE_DEVICES"] = self.settings.arbor_config.inference.gpu_ids
-        n_gpus = self.settings.arbor_config.inference.gpu_ids.count(",") + 1
-        # command = f"vllm serve {model} --port {port} --gpu-memory-utilization 0.9 --tensor-parallel-size {n_gpus} --max_model_len 8192 --enable_prefix_caching"
-        command = f"python -m sglang_router.launch_server --model-path {model} --dp-size {n_gpus} --port {port} --host 0.0.0.0 --disable-radix-cache"
-        print(f"Running command: {command}")
-
-        # We will manually stream & capture logs.
-        process = subprocess.Popen(
-            command.replace("\\\n", " ").replace("\\", " ").split(),
-            text=True,
-            stdout=subprocess.PIPE,  # We'll read from pipe
-            stderr=subprocess.STDOUT,  # Merge stderr into stdout
-            env=my_env,
-        )
-
-        # A threading.Event to control printing after the server is ready.
-        # This will store *all* lines (both before and after readiness).
-        print(f"SGLang server process started with PID {process.pid}.")
-        stop_printing_event = threading.Event()
-        logs_buffer = []
-
-        def _tail_process(proc, buffer, stop_event):
-            while True:
-                line = proc.stdout.readline()
-                if not line and proc.poll() is not None:
-                    # Process ended and no new line
-                    break
-                if line:
-                    buffer.append(line)
-                    # Print only if stop_event is not set
-                    if not stop_event.is_set():
-                        print(f"[SGLang LOG] {line}", end="")
-
-        # Start a background thread to read from the process continuously
-        thread = threading.Thread(
-            target=_tail_process,
-            args=(process, logs_buffer, stop_printing_event),
-            daemon=True,
-        )
-        thread.start()
-
-        # Wait until the server is ready (or times out)
-        base_url = f"http://localhost:{port}"
-        try:
-            wait_for_server(base_url, timeout=timeout)
-        except TimeoutError:
-            # If the server doesn't come up, we might want to kill it:
-            process.kill()
-            raise
-
-        # Once server is ready, we tell the thread to stop printing further lines.
-        stop_printing_event.set()
-
-        # A convenience getter so the caller can see all logs so far (and future).
-        def get_logs() -> str:
-            # Join them all into a single string, or you might return a list
-            return "".join(logs_buffer)
-
-        # Let the user know server is up
-        print(f"Server ready on random port {port}!")
+        retries = 0
+        while retries < max_retries:
+            try:
+                print(
+                    f"Attempt {retries + 1} of {max_retries} to launch server for model {model}"
+                )
+                print(
+                    f"Grabbing a free port to launch an SGLang server for model {model}"
+                )
+                router_port = get_free_port()
+                dp_worker_base_port = get_free_port()
+                worker_urls_port = get_free_port()  # Get a port for worker URLs
+
+                timeout = launch_kwargs.get("timeout", 1800)
+                my_env = os.environ.copy()
+                my_env["CUDA_VISIBLE_DEVICES"] = (
+                    self.settings.arbor_config.inference.gpu_ids
+                )
+                n_gpus = self.settings.arbor_config.inference.gpu_ids.count(",") + 1
+                command = f"python -m arbor.server.services.inference.sgl_router_launch_server --model-path {model} --dp-size {n_gpus} --port {router_port} --host 0.0.0.0 --disable-radix-cache --router-dp-worker-base-port {dp_worker_base_port} --worker-urls-port {worker_urls_port}"
+                print(f"Running command: {command}")
+                if launch_kwargs.get("max_context_length"):
+                    command += (
+                        f" --context-length {launch_kwargs['max_context_length']}"
+                    )
 
-        self.launch_kwargs["api_base"] = f"http://localhost:{port}/v1"
-        self.launch_kwargs["api_key"] = "local"
-        self.get_logs = get_logs
-        self.process = process
-        self.thread = thread
-        self.current_model = model
+                # We will manually stream & capture logs.
+                process = subprocess.Popen(
+                    command.replace("\\\n", " ").replace("\\", " ").split(),
+                    text=True,
+                    stdout=subprocess.PIPE,  # We'll read from pipe
+                    stderr=subprocess.STDOUT,  # Merge stderr into stdout
+                    env=my_env,
+                )
+
+                # A threading.Event to control printing after the server is ready.
+                # This will store *all* lines (both before and after readiness).
+                print(f"SGLang server process started with PID {process.pid}.")
+                stop_printing_event = threading.Event()
+                logs_buffer = []
+
+                def _tail_process(proc, buffer, stop_event):
+                    while True:
+                        line = proc.stdout.readline()
+                        if not line and proc.poll() is not None:
+                            # Process ended and no new line
+                            break
+                        if line:
+                            buffer.append(line)
+                            # Print only if stop_event is not set
+                            if not stop_event.is_set():
+                                print(f"[SGLang LOG] {line}", end="")
+
+                # Start a background thread to read from the process continuously
+                thread = threading.Thread(
+                    target=_tail_process,
+                    args=(process, logs_buffer, stop_printing_event),
+                    daemon=True,
+                )
+                thread.start()
+
+                # Get worker URLs before waiting for server
+                try:
+                    worker_urls = get_worker_urls(worker_urls_port)
+                    print(f"Received worker URLs: {worker_urls}")
+                    self.worker_urls = worker_urls
+                except TimeoutError as e:
+                    raise Exception(f"Failed to get worker URLs: {e}")
+
+                # Wait until the server is ready (or times out)
+                base_url = f"http://localhost:{router_port}"
+                try:
+                    wait_for_server(base_url, timeout=timeout)
+                except TimeoutError:
+                    # If the server doesn't come up, we might want to kill it:
+                    process.kill()
+                    raise
+
+                # Once server is ready, we tell the thread to stop printing further lines.
+                stop_printing_event.set()
+
+                # A convenience getter so the caller can see all logs so far (and future).
+                def get_logs() -> str:
+                    # Join them all into a single string, or you might return a list
+                    return "".join(logs_buffer)
+
+                # Let the user know server is up
+                print(f"Server ready on random port {router_port}!")
+
+                self.launch_kwargs["api_base"] = f"http://localhost:{router_port}/v1"
+                self.launch_kwargs["api_key"] = "local"
+                self.get_logs = get_logs
+                self.process = process
+                self.thread = thread
+                self.current_model = model
+
+                # If we get here, the launch was successful
+                return
+
+            except Exception as e:
+                retries += 1
+                print(
+                    f"Failed to launch server (attempt {retries} of {max_retries}): {str(e)}"
+                )
+                # Clean up any failed processes
+                if "process" in locals():
+                    try:
+                        process.kill()
+                    except:
+                        pass
+                if retries == max_retries:
+                    raise Exception(
+                        f"Failed to launch server after {max_retries} attempts"
+                    ) from e
+                # Wait a bit before retrying
+                time.sleep(min(2**retries, 30))  # Exponential backoff, max 30 seconds
 
     def kill(self):
         from sglang.utils import terminate_process
@@ -184,7 +237,7 @@ class InferenceManager:
         print(f"Running inference for model {model}")
         # Monkeypatch:
         if model != self.current_model:
-            print(f"MONKEYPATCH: Model changed from {model} to {self.current_model}")
+            print(f"Model changed from {model} to {self.current_model}")
             model = self.current_model
             request_json["model"] = model
 
@@ -214,6 +267,12 @@ class InferenceManager:
                 await self._session.close()
                 self._session = None
             return None
+        except json.decoder.JSONDecodeError:
+            print(f"JSON Decode Error during inference: {content}")
+            return {
+                "error": "JSON Decode Error",
+                "content": content if content else "Content is null",
+            }
         except Exception as e:
             print(f"Error during inference: {e}")
             raise
@@ -239,8 +298,10 @@ class InferenceManager:
         self.inference_count = 0
 
         tik = time.time()
-        self.kill()
-        print("Just killed server")
+        # self.kill()
+        # print("Just killed server")
+        # time.sleep(5)
+
         # Check that output directory exists and was created successfully
         print(f"Checking that output directory {output_dir} exists")
         if not os.path.exists(output_dir):
@@ -248,8 +309,27 @@ class InferenceManager:
                 f"Failed to save model - output directory {output_dir} does not exist"
             )
 
-        print("Launching new server")
-        self.launch(output_dir, self.launch_kwargs)
+        print("Directly updating weights from disk")
+        for worker_url in self.worker_urls:
+            print(f"Updating weights from disk for worker {worker_url}")
+            try:
+                response = requests.post(
+                    f"{worker_url}/update_weights_from_disk",
+                    json={"model_path": output_dir},
+                )
+                response_json = response.json()
+                print(f"Response from update_weights_from_disk: {response_json}")
+                # TODO: Check that the response is successful
+            except Exception as e:
+                print(f"Error during update_weights_from_disk: {e}")
+                print(f"Full error during update_weights_from_disk: {str(e)}")
+                if hasattr(e, "response") and e.response is not None:
+                    print(f"Response status code: {e.response.status_code}")
+                    print(f"Response text: {e.response.text}")
+        self.current_model = output_dir
+
+        # print("Launching new server")
+        # self.launch(output_dir, self.launch_kwargs)
         tok = time.time()
         self.restarting = False
         print(f"Time taken to update model: {tok - tik} seconds")
@@ -297,3 +377,28 @@ def wait_for_server(base_url: str, timeout: int = None) -> None:
         except requests.exceptions.RequestException:
             # Server not up yet, wait and retry
             time.sleep(1)
+
+
+def get_worker_urls(zmq_port: int, timeout: float = 30.0) -> list:
+    print(f"Attempting to get worker URLs on port {zmq_port} with timeout {timeout}s")
+    context = zmq.Context()
+    socket = context.socket(zmq.SUB)
+    socket.connect(f"tcp://localhost:{zmq_port}")
+    socket.setsockopt_string(zmq.SUBSCRIBE, "")  # Subscribe to all messages
+
+    # Set a timeout for receiving
+    socket.setsockopt(zmq.RCVTIMEO, int(timeout * 1000))
+
+    try:
+        print("Waiting for worker URLs message...")
+        message = socket.recv_json()
+        print(f"Received message: {message}")
+        if message.get("type") == "worker_urls":
+            return message["urls"]
+        else:
+            raise ValueError(f"Unexpected message type: {message.get('type')}")
+    except zmq.error.Again:
+        raise TimeoutError(f"Timeout waiting for worker URLs on port {zmq_port}")
+    finally:
+        socket.close()
+        context.term()

arbor/server/services/scripts/grpo_training.py

@@ -14,7 +14,7 @@ from typing import Any, List, Optional, Union
 import torch
 import zmq
 from accelerate import Accelerator
-from accelerate.utils import gather
+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
@@ -23,6 +23,7 @@ from transformers import (
     PreTrainedTokenizerBase,
     Trainer,
     TrainerCallback,
+    is_wandb_available,
 )
 from trl import GRPOConfig, GRPOTrainer
 from trl.data_utils import maybe_apply_chat_template
@@ -32,6 +33,9 @@ from arbor.server.services.comms.comms import (
     ArborServerCommsHandler,
 )
 
+if is_wandb_available():
+    import wandb
+
 last_step_time = None
 last_queue_pop_time = None
 
@@ -65,8 +69,9 @@ class ArborGRPOTrainer(GRPOTrainer):
         ] = (None, None),
         peft_config: Optional["PeftConfig"] = None,
         comms_handler: Optional[ArborScriptCommsHandler] = None,
-        update_interval: Optional[int] = 5,
         lora: Optional[bool] = False,
+        # We do nothing with max_context_length right now
+        max_context_length: Optional[int] = None,
         **kwargs,
     ):
 
@@ -85,12 +90,12 @@ class ArborGRPOTrainer(GRPOTrainer):
         self.peft_config = peft_config
         self.scale_rewards = scale_rewards
         self.comms_handler = comms_handler
-        self.update_interval = update_interval
 
     def _generate_and_score_completions(
         self, batch: List[dict[str, Any]]
     ) -> dict[str, Union[torch.Tensor, Any]]:
         device = self.accelerator.device
+        mode = "train" if self.model.training else "eval"
 
         # Process prompts and completions
         prompt_completion_texts = []
@@ -106,12 +111,12 @@ class ArborGRPOTrainer(GRPOTrainer):
             )
 
         # Tokenize prompts
-        prompt_texts = [
+        prompts_text = [
             prompt_completion_text["prompt"]
             for prompt_completion_text in prompt_completion_texts
         ]
         prompt_inputs = self.processing_class(
-            prompt_texts,
+            prompts_text,
             return_tensors="pt",
             padding=True,
             padding_side="left",
@@ -124,12 +129,12 @@ class ArborGRPOTrainer(GRPOTrainer):
         )
 
         # Tokenize completions
-        completion_texts = [
+        completions_text = [
             prompt_completion_text["completion"]
             for prompt_completion_text in prompt_completion_texts
         ]
         completion_ids = self.processing_class(
-            completion_texts,
+            completions_text,
             return_tensors="pt",
             padding=True,
             add_special_tokens=False,
@@ -156,6 +161,30 @@ class ArborGRPOTrainer(GRPOTrainer):
         #     self._move_model_to_vllm()
         #     self._last_loaded_step = self.state.global_step
 
+        prompt_ids = broadcast_object_list(prompt_ids)
+        prompt_mask = broadcast_object_list(prompt_mask)
+        completion_ids = broadcast_object_list(completion_ids)
+        completion_mask = broadcast_object_list(completion_mask)
+
+        process_slice = slice(
+            self.accelerator.process_index * len(batch),
+            (self.accelerator.process_index + 1) * len(batch),
+        )
+
+        prompt_ids = prompt_ids[process_slice]
+        prompt_mask = prompt_mask[process_slice]
+        completion_ids = completion_ids[process_slice]
+        completion_mask = completion_mask[process_slice]
+
+        is_eos = completion_ids == self.processing_class.eos_token_id
+
+        # If mask_truncated_completions is enabled, zero out truncated completions in completion_mask
+        if self.mask_truncated_completions:
+            truncated_completions = ~is_eos.any(dim=1)
+            completion_mask = (
+                completion_mask * (~truncated_completions).unsqueeze(1).int()
+            )
+
         prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1)
         attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)  # (B, P+C)
 
@@ -164,34 +193,29 @@ class ArborGRPOTrainer(GRPOTrainer):
         )
 
         logits_to_keep = completion_ids.size(1)
+        batch_size = (
+            self.args.per_device_train_batch_size
+            if mode == "train"
+            else self.args.per_device_eval_batch_size
+        )
 
         with torch.no_grad():
             # When using num_iterations == 1, old_per_token_logps == per_token_logps, so we can skip it's
             # computation here, and use per_token_logps.detach() instead.
-            if self.num_iterations > 1:
+            if (
+                self.num_iterations > 1
+                or self.args.steps_per_generation
+                > self.args.gradient_accumulation_steps
+            ):
                 old_per_token_logps = self._get_per_token_logps(
-                    self.model, prompt_completion_ids, attention_mask, logits_to_keep
-                )
-            else:
-                old_per_token_logps = None
-
-            if self.beta == 0.0:
-                ref_per_token_logps = None
-            elif self.ref_model is not None:
-                ref_per_token_logps = self._get_per_token_logps(
-                    self.ref_model,
+                    self.model,
                     prompt_completion_ids,
                     attention_mask,
                     logits_to_keep,
+                    batch_size,
                 )
             else:
-                with self.accelerator.unwrap_model(self.model).disable_adapter():
-                    ref_per_token_logps = self._get_per_token_logps(
-                        self.model,
-                        prompt_completion_ids,
-                        attention_mask,
-                        logits_to_keep,
-                    )
+                old_per_token_logps = None
 
         rewards = torch.tensor(
             [example["reward"] for example in batch], dtype=torch.float32
@@ -219,7 +243,56 @@ class ArborGRPOTrainer(GRPOTrainer):
         )
         advantages = advantages[process_slice]
 
-        ## Logged Metrics Removed Here
+        # Log the metrics
+        if mode == "train":
+            self.state.num_input_tokens_seen += (
+                self.accelerator.gather_for_metrics(attention_mask.sum()).sum().item()
+            )
+        self._metrics[mode]["num_tokens"] = [self.state.num_input_tokens_seen]
+
+        # log completion lengths, mean, min, max
+        agg_completion_mask = self.accelerator.gather_for_metrics(
+            completion_mask.sum(1)
+        )
+        self._metrics[mode]["completions/mean_length"].append(
+            agg_completion_mask.float().mean().item()
+        )
+        self._metrics[mode]["completions/min_length"].append(
+            agg_completion_mask.float().min().item()
+        )
+        self._metrics[mode]["completions/max_length"].append(
+            agg_completion_mask.float().max().item()
+        )
+
+        # identify sequences that terminated with EOS and log their lengths
+        agg_terminated_with_eos = self.accelerator.gather_for_metrics(is_eos.any(dim=1))
+        term_completion_mask = agg_completion_mask[agg_terminated_with_eos]
+        clipped_completions_ratio = 1 - len(term_completion_mask) / len(
+            agg_completion_mask
+        )
+        self._metrics[mode]["completions/clipped_ratio"].append(
+            clipped_completions_ratio
+        )
+        if len(term_completion_mask) == 0:
+            # edge case where no completed sequences are found
+            term_completion_mask = torch.zeros(1, device=device)
+        self._metrics[mode]["completions/mean_terminated_length"].append(
+            term_completion_mask.float().mean().item()
+        )
+        self._metrics[mode]["completions/min_terminated_length"].append(
+            term_completion_mask.float().min().item()
+        )
+        self._metrics[mode]["completions/max_terminated_length"].append(
+            term_completion_mask.float().max().item()
+        )
+
+        # Calculate mean reward
+        self._metrics[mode]["reward"].append(mean_grouped_rewards.mean().item())
+        self._metrics[mode]["reward_std"].append(std_grouped_rewards.mean().item())
+
+        # Log prompt and completion texts
+        self._textual_logs["prompt"].extend(gather_object(prompts_text))
+        self._textual_logs["completion"].extend(gather_object(completions_text))
 
         return {
             "prompt_ids": prompt_ids,
@@ -227,7 +300,6 @@ class ArborGRPOTrainer(GRPOTrainer):
             "completion_ids": completion_ids,
             "completion_mask": completion_mask,
             "old_per_token_logps": old_per_token_logps,
-            "ref_per_token_logps": ref_per_token_logps,
             "advantages": advantages,
         }
 
@@ -326,15 +398,10 @@ class CommandMonitor:
                     print(
                         f"[Training Script] Instructed to save model at {self.trainer.args.output_dir}"
                     )
-                    # Wait until data queue is empty before saving
-
                     while (
                         time_since_last_step() <= 10
                         or get_time_since_last_queue_pop() <= 10
                     ):
-                        # print(
-                        # f"Waiting for data queue to empty...{self.comms_handler.get_data_queue_size()}"
-                        # )
                         print(f"Waiting for steps to finish")
                         print(
                             f"Time since last step: {time_since_last_step():.1f} (needs to be >= 10)"
@@ -342,15 +409,12 @@ class CommandMonitor:
                         print(
                             f"Time since last queue pop: {get_time_since_last_queue_pop():.1f} (needs to be >= 10)"
                         )
-                        time.sleep(5)  # Small delay to prevent busy waiting)
+                        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,
@@ -373,6 +437,56 @@ class CommandMonitor:
                             "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')}/"
+                        )
+                    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')}/",
+                        }
+                    )
+
         except Exception as e:
             print(e)
             self.comms_handler.send_status({"status": "error", "error": str(e)})
@@ -385,13 +499,15 @@ class CommandMonitor:
             for broadcast in self.comms_handler.receive_broadcast():
                 print(f"!!!Received broadcast: {broadcast}")
                 if broadcast.get("message") == "terminate":
-                    self.trainer.control.should_training_stop = True
-                    self.comms_handler.send_status(
-                        {
-                            "status": "Received termination command",
-                            "process_id": self.trainer.accelerator.process_index,
-                        }
-                    )
+                    # self.trainer.control.should_training_stop = True
+                    # self.comms_handler.send_status(
+                    #     {
+                    #         "status": "Received termination command",
+                    #         "process_id": self.trainer.accelerator.process_index,
+                    #     }
+                    # )
+                    if self.trainer.accelerator.is_main_process:
+                        self.trainer.accelerator.end_training()
         except Exception as e:
             self.comms_handler.send_status({"status": "error", "error": str(e)})
 

{arbor_ai-0.1.12.dist-info → arbor_ai-0.1.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arbor-ai
-Version: 0.1.12
+Version: 0.1.14
 Summary: A framework for fine-tuning and managing language models
 Author-email: Noah Ziems <nziems2@nd.edu>
 Project-URL: Homepage, https://github.com/Ziems/arbor
@@ -23,6 +23,7 @@ Requires-Dist: pyzmq>=26.4.0
 Requires-Dist: pyyaml>=6.0.2
 Requires-Dist: sglang[all]>=0.4.5.post3
 Requires-Dist: sglang-router
+Requires-Dist: wandb
 Dynamic: license-file
 
 <p align="center">

{arbor_ai-0.1.12.dist-info → arbor_ai-0.1.14.dist-info}/RECORD

@@ -5,10 +5,10 @@ arbor/client/api.py,sha256=86bgHuGM_AvI1Uhic_QaCnpF4VFqXie9ZzxmbTXUPpQ,19
 arbor/server/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 arbor/server/main.py,sha256=tY4Vlaaj4oq1FTGYOkbFMGF0quLEeR-VBaKaXhQ5mEE,382
 arbor/server/api/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-arbor/server/api/models/schemas.py,sha256=s_G8sSb05FjkKEqpKpLlqaEd8NysJddHibRHhcnrKIk,5594
+arbor/server/api/models/schemas.py,sha256=KCHav1nPFbQEynrcO-MObhRmoOrdFvfGuVogApynOCA,6210
 arbor/server/api/routes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 arbor/server/api/routes/files.py,sha256=DQC_ogH5zlzhHZSAA4Cj5wzK07XBIBVs2Po91W9rcDY,1835
-arbor/server/api/routes/grpo.py,sha256=VuEvSOwwrHegn9qM-1nbHFmmUnnC_BMwnIHsfIdiJyI,1877
+arbor/server/api/routes/grpo.py,sha256=AbQ_BHgk-Om5U0qSt_FeJfyBJ0vItpfrnCNtJgD6p5k,2245
 arbor/server/api/routes/inference.py,sha256=Zy4ciN6vdRgu0-sFFnEeTZB-4XnLjEDH-atU7roIKSs,1668
 arbor/server/api/routes/jobs.py,sha256=BNdaSYUBJX6xSd6Pj6qx1DQJiZ5EKVxxbXDbEkfkCpw,3634
 arbor/server/core/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
@@ -17,18 +17,20 @@ arbor/server/core/logging.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,
 arbor/server/services/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 arbor/server/services/dependencies.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 arbor/server/services/file_manager.py,sha256=Z9z4A4EzvPauid_DBfpim401DDtuJy_TbX4twTWDJWI,12119
-arbor/server/services/grpo_manager.py,sha256=TAU2BMHgbCgiAvKNVd2Y8N20SR4qEms3lChA4Z0ZzyY,13777
-arbor/server/services/inference_manager.py,sha256=q4RVUqh1snGfW-AADkCqW8hC5x3WAZNe0jwXKOY5joU,10685
+arbor/server/services/grpo_manager.py,sha256=-_0xjENvIrOAtHACkFPMYox9YAeckHbpX2FkrmKrWuU,15448
+arbor/server/services/inference_manager.py,sha256=Ju39_7EWySzAAk7ftz-AzSNBEo0tlayloPVS0XRAp8E,15304
 arbor/server/services/job_manager.py,sha256=m_d4UPwN_82f7t7K443DaFpFoyv7JZSZKml8tawt1Bk,2186
 arbor/server/services/training_manager.py,sha256=oQdhpfxdgp_lCTb_lxhvjupdLrcg6HL3TEbct_q9F6I,21065
 arbor/server/services/comms/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 arbor/server/services/comms/comms.py,sha256=3KN3mzwPvfW2_L5hq02JdAk6yOMyhY0_pBz-DDr5A3o,7694
-arbor/server/services/scripts/grpo_training.py,sha256=Q9jwnbRdXAv_jVgrChLX6IiB3BLZU1F3BP6mBV0DVik,20889
+arbor/server/services/inference/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+arbor/server/services/inference/sgl_router_launch_server.py,sha256=eqTW6nDqqoRMISHfv5ScBCrolqLBp9zyxPXqHUlP6uo,6988
+arbor/server/services/scripts/grpo_training.py,sha256=eMT5cIMolAzhukANH1WRmPdxIkvLbsbrggdGFCMGMHc,26474
 arbor/server/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 arbor/server/utils/helpers.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-arbor_ai-0.1.12.dist-info/licenses/LICENSE,sha256=5vFGrbOFeXXM83JV9o16w7ohH4WLeu3-57GocJSz8ow,1067
-arbor_ai-0.1.12.dist-info/METADATA,sha256=upqnB_F9JDLytHm4AFrDnvPaOHdj8XiBCdrlam0rgRc,2413
-arbor_ai-0.1.12.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
-arbor_ai-0.1.12.dist-info/entry_points.txt,sha256=PGBX-MfNwfIl8UPFgsX3gjtXLqSogRhOktKMpZUysD0,40
-arbor_ai-0.1.12.dist-info/top_level.txt,sha256=jzWdp3BRYqvZDMFsPajrcftvvlluzVDErkD8IMRfhYs,6
-arbor_ai-0.1.12.dist-info/RECORD,,
+arbor_ai-0.1.14.dist-info/licenses/LICENSE,sha256=5vFGrbOFeXXM83JV9o16w7ohH4WLeu3-57GocJSz8ow,1067
+arbor_ai-0.1.14.dist-info/METADATA,sha256=vw8RnMPdGi36ji4rpjAldkOuCbxxjV4MFVi6yW-0kas,2434
+arbor_ai-0.1.14.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+arbor_ai-0.1.14.dist-info/entry_points.txt,sha256=PGBX-MfNwfIl8UPFgsX3gjtXLqSogRhOktKMpZUysD0,40
+arbor_ai-0.1.14.dist-info/top_level.txt,sha256=jzWdp3BRYqvZDMFsPajrcftvvlluzVDErkD8IMRfhYs,6
+arbor_ai-0.1.14.dist-info/RECORD,,

{arbor_ai-0.1.12.dist-info → arbor_ai-0.1.14.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.4.0)
+Generator: setuptools (80.7.1)
 Root-Is-Purelib: true
 Tag: py3-none-any