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

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

@@ -13,6 +13,7 @@ from typing import Any, Dict, Optional
 
 import aiohttp
 import requests
+import zmq
 
 from arbor.server.core.config import Settings
 
@@ -28,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)
@@ -74,15 +76,17 @@ class InferenceManager:
                 print(
                     f"Grabbing a free port to launch an SGLang server for model {model}"
                 )
-                port = get_free_port()
+                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"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"
+                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 += (
@@ -124,8 +128,16 @@ class InferenceManager:
                 )
                 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:{port}"
+                base_url = f"http://localhost:{router_port}"
                 try:
                     wait_for_server(base_url, timeout=timeout)
                 except TimeoutError:
@@ -142,9 +154,9 @@ class InferenceManager:
                     return "".join(logs_buffer)
 
                 # Let the user know server is up
-                print(f"Server ready on random port {port}!")
+                print(f"Server ready on random port {router_port}!")
 
-                self.launch_kwargs["api_base"] = f"http://localhost:{port}/v1"
+                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
@@ -286,9 +298,10 @@ class InferenceManager:
         self.inference_count = 0
 
         tik = time.time()
-        self.kill()
-        print("Just killed server")
-        time.sleep(5)
+        # 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):
@@ -296,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")
@@ -345,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_ai-0.1.13.dist-info → arbor_ai-0.1.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arbor-ai
-Version: 0.1.13
+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
@@ -15,7 +15,7 @@ Requires-Dist: python-multipart
 Requires-Dist: pydantic-settings
 Requires-Dist: torch
 Requires-Dist: transformers
-Requires-Dist: trl==0.17.0
+Requires-Dist: trl
 Requires-Dist: peft
 Requires-Dist: ray>=2.9
 Requires-Dist: setuptools<77.0.0,>=76.0.0

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

@@ -18,17 +18,19 @@ arbor/server/services/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3h
 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=-_0xjENvIrOAtHACkFPMYox9YAeckHbpX2FkrmKrWuU,15448
-arbor/server/services/inference_manager.py,sha256=NcsUI-pgf3cRhU6P3xlPx0dxhvgYrfGZkEEGORcHcis,12833
+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/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.13.dist-info/licenses/LICENSE,sha256=5vFGrbOFeXXM83JV9o16w7ohH4WLeu3-57GocJSz8ow,1067
-arbor_ai-0.1.13.dist-info/METADATA,sha256=c0yScMpCiWYSFqVLjgk5TrRBuAVJK3aTBl0z0IPZ_8Y,2442
-arbor_ai-0.1.13.dist-info/WHEEL,sha256=QZxptf4Y1BKFRCEDxD4h2V0mBFQOVFLFEpvxHmIs52A,91
-arbor_ai-0.1.13.dist-info/entry_points.txt,sha256=PGBX-MfNwfIl8UPFgsX3gjtXLqSogRhOktKMpZUysD0,40
-arbor_ai-0.1.13.dist-info/top_level.txt,sha256=jzWdp3BRYqvZDMFsPajrcftvvlluzVDErkD8IMRfhYs,6
-arbor_ai-0.1.13.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.13.dist-info → arbor_ai-0.1.14.dist-info}/WHEEL

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