arbor-ai 0.1.12__tar.gz → 0.1.14__tar.gz

{arbor_ai-0.1.12/arbor_ai.egg-info → arbor_ai-0.1.14}/PKG-INFO

@@ -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 → arbor_ai-0.1.14}/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_ai-0.1.12 → arbor_ai-0.1.14}/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_ai-0.1.12 → arbor_ai-0.1.14}/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_ai-0.1.14/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()