arbor-ai 0.1.5__py3-none-any.whl → 0.1.6__py3-none-any.whl

arbor/server/services/grpo_manager.py

@@ -0,0 +1,310 @@
+import json
+import os
+import random
+import signal
+import string
+import subprocess
+import sys
+import threading
+import time
+from datetime import datetime
+from pathlib import Path
+from typing import Optional
+
+from arbor.server.api.models.schemas import 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
+
+
+class GRPOManager:
+    def __init__(self, settings: Settings):
+        self.settings = settings
+        self.training_process = None
+        self.current_model = None
+        self.train_kwargs = None
+        self.server_comms_handler = None
+        self.status_thread = None
+        self.model_saved_and_reload_requested = False
+
+        self.data_count = 0
+        self.last_inference_update = 0
+        # Set up signal handler
+        signal.signal(signal.SIGINT, self._signal_handler)
+        signal.signal(signal.SIGTERM, self._signal_handler)
+
+    def _signal_handler(self, signum, frame):
+        """Handle keyboard interrupt (SIGINT) gracefully."""
+        print("\nReceived keyboard interrupt. Shutting down gracefully...")
+        self.terminate(None)
+        sys.exit(0)
+
+    def make_output_dir(
+        self, model_name: str, run_suffix: Optional[str] = None
+    ) -> tuple[str, str]:
+        """Create a unique output directory name for the training run."""
+        model_name = model_name.split("/")[-1].lower()
+        suffix = (
+            run_suffix
+            if run_suffix
+            else "".join(random.choices(string.ascii_letters + string.digits, k=6))
+        )
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        name = f"grpo:{model_name}:{suffix}:{timestamp}"
+        return name, str(Path(self.settings.STORAGE_PATH).resolve() / "models" / name)
+
+    def find_training_args(self, request: GRPOConfigRequest) -> dict:
+        """Process the config request and return training arguments."""
+        name, output_dir = self.make_output_dir(request.model, request.suffix)
+
+        # Here are defaults for training. We can adjust them if we disagree w the huggingface defaults
+        default_train_kwargs = {
+            "output_dir": output_dir,
+        }
+
+        train_kwargs = request.model_dump(exclude_unset=True)
+        return {**default_train_kwargs, **(train_kwargs or {})}
+
+    def process_training_args(self, train_kwargs: dict) -> tuple[dict, dict]:
+        # NOTE: These also need to be in the GRPOConfigRequest
+        trl_keys = [
+            "output_dir",
+            "temperature",
+            "beta",
+            "num_iterations",
+            "num_generations",
+            "per_device_train_batch_size",
+            "learning_rate",
+            "gradient_accumulation_steps",
+            "gradient_checkpointing",
+            "lr_scheduler_type",
+            "max_prompt_length",
+            "max_completion_length",
+            "gradient_checkpointing_kwargs",
+            "bf16",
+            "scale_rewards",
+            "max_grad_norm",
+        ]
+        trl_train_kwargs = {
+            key: train_kwargs[key] for key in trl_keys if key in train_kwargs
+        }
+
+        arbor_keys = ["update_interval", "lora"]
+        arbor_train_kwargs = {
+            key: train_kwargs[key] for key in arbor_keys if key in train_kwargs
+        }
+
+        return trl_train_kwargs, arbor_train_kwargs
+
+    def initialize(
+        self, request: GRPOConfigRequest, inference_manager: InferenceManager
+    ):
+        """Initialize the training process with ZMQ-based communication."""
+        self.train_kwargs = self.find_training_args(request)
+
+        trl_train_kwargs, arbor_train_kwargs = self.process_training_args(
+            self.train_kwargs
+        )
+
+        self.current_model = request.model
+
+        # Initialize ZMQ socket manager - no need for connection acceptance thread anymore
+        self.server_comms_handler = ArborServerCommsHandler()
+
+        script_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "scripts")
+        script_path = os.path.join(script_dir, "grpo_training.py")
+
+        # Start the training process with ZMQ ports
+        my_env = os.environ.copy()
+        my_env["CUDA_VISIBLE_DEVICES"] = self.settings.arbor_config.training.gpu_ids
+
+        num_processes = self.settings.arbor_config.training.gpu_ids.count(",") + 1
+
+        params = [
+            "python",
+            "-m",
+            "accelerate.commands.launch",
+            "--num_processes",
+            str(num_processes),
+        ]
+        if self.settings.arbor_config.training.accelerate_config:
+            params.extend(
+                [
+                    "--config_file",
+                    self.settings.arbor_config.training.accelerate_config,
+                ]
+            )
+        params.extend(
+            [
+                script_path,
+                # Comms args
+                "--host",
+                self.server_comms_handler.host,
+                "--command_port",
+                str(self.server_comms_handler.command_port),
+                "--status_port",
+                str(self.server_comms_handler.status_port),
+                "--data_port",
+                str(self.server_comms_handler.data_port),
+                "--broadcast_port",
+                str(self.server_comms_handler.broadcast_port),
+                "--handshake_port",
+                str(self.server_comms_handler.handshake_port),
+                # Training args
+                "--model",
+                self.current_model,
+                "--trl_train_kwargs",
+                json.dumps(trl_train_kwargs),
+                "--arbor_train_kwargs",
+                json.dumps(arbor_train_kwargs),
+            ]
+        )
+        print(f"Running following command\n: {' '.join(params)}")
+
+        self.training_process = subprocess.Popen(
+            params,
+            text=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            env=my_env,
+        )
+
+        # A threading.Event to control printing after the server is ready.
+        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"[GRPO LOG] {line}", end="")
+
+        # Start a background thread to read from the process continuously
+        thread = threading.Thread(
+            target=_tail_process,
+            args=(self.training_process, logs_buffer, stop_printing_event),
+            daemon=True,
+        )
+        thread.start()
+
+        # Start status handling thread
+        self.status_thread = threading.Thread(
+            target=self._handle_status_updates, args=(inference_manager,), daemon=True
+        )
+        self.status_thread.start()
+        self.server_comms_handler.wait_for_clients(num_processes)
+
+        # Launch the inference server
+        print("Launching inference server...")
+        inference_manager.launch(self.current_model)
+
+    def _handle_status_updates(self, inference_manager: InferenceManager):
+        """Handle status updates from training process using ZMQ SUB socket"""
+        print("Starting status update handler...")
+        try:
+
+            for status in self.server_comms_handler.receive_status():
+                print(f"Received status update: {status}")
+                if status["status"] == "model_saved":
+                    print("Updating inference model...")
+                    # There is a case where this status is sent multiple times
+                    # We need to make sure we only update the model once
+                    if self._should_update_model():
+                        inference_manager.update_model(status["output_dir"])
+                        # self.last_inference_update = self.data_count
+                        self.model_saved_and_reload_requested = False
+                        self.current_model = status["output_dir"]
+                        print("Model update complete")
+                elif status["status"] == "error":
+                    print(f"Training error: {status.get('error', 'Unknown error')}")
+                elif status["status"] == "terminated":
+                    print("Training process terminated")
+                    break
+        except Exception as e:
+            print(f"Error in status update handler: {e}")
+
+    def grpo_step(
+        self, request: GRPORequest, inference_manager: InferenceManager
+    ) -> str:
+        while inference_manager.is_server_restarting():
+            print("Inferece manager restarting, waiting for GRPO step")
+            time.sleep(5)
+
+        while self._should_update_model():
+            print(
+                f"Waiting for model update. Data count: {self.data_count}, Last inference update: {self.last_inference_update}"
+            )
+            time.sleep(5)
+
+        try:
+            # Send the batch to the training process
+            self.server_comms_handler.send_data(request.batch)
+            self.data_count += 1
+        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
+
+    def update_model(self, request, inference_manager: InferenceManager):
+        # THIS IS HACKY AND NEEDS TO BE FIXED BEFORE RELEASE
+        inference_manager.restarting = True
+        self.model_saved_and_reload_requested = True
+        self.server_comms_handler.send_command({"command": "save_model"})
+        while self.model_saved_and_reload_requested:
+            print(
+                "Waiting for model to be saved and reloaded... This usually takes 20-30 seconds"
+            )
+            time.sleep(5)
+        return self.current_model
+
+    def terminate(self, inference_manager: InferenceManager):
+        """Clean up resources and save the final model."""
+        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"})
+
+            # Wait for training process to finish
+            if self.training_process:
+                self.training_process.wait(timeout=30)
+
+        except Exception as e:
+            print(f"Error during termination: {e}")
+        finally:
+            # Clean up ZMQ connections
+            if self.server_comms_handler:
+                self.server_comms_handler.close()
+
+            if self.train_kwargs and "output_dir" in self.train_kwargs:
+                print(
+                    f"Training completed. Model saved to {self.train_kwargs['output_dir']}"
+                )
+                if not os.path.exists(self.train_kwargs["output_dir"]):
+                    print(
+                        f"Warning: Output directory {self.train_kwargs['output_dir']} does not exist"
+                    )
+                return self.train_kwargs["output_dir"]
+            else:
+                print("Training terminated, no output directory specified")
+                return None
+
+    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_manager.py

@@ -0,0 +1,275 @@
+import os
+import signal
+import socket
+import subprocess
+import sys
+import threading
+import time
+from datetime import datetime
+from typing import Any, Dict, Optional
+
+import requests
+
+from arbor.server.core.config import Settings
+
+
+class InferenceManager:
+    def __init__(self, settings: Settings):
+        self.settings = settings
+        self.process = None
+        self.launch_kwargs = {}
+        self.last_activity = None
+        self.restarting = False
+        self._shutting_down = False
+        self.current_model = None
+        self.inference_count = 0
+        # Set up signal handler for graceful shutdown
+        signal.signal(signal.SIGINT, self._signal_handler)
+        signal.signal(signal.SIGTERM, self._signal_handler)
+
+    def _signal_handler(self, signum, frame):
+        if self._shutting_down:
+            print("\nForced exit during cleanup...")
+            os._exit(1)
+
+        print("\nReceived signal to terminate. Cleaning up...")
+        self._shutting_down = True
+        self.kill()
+        sys.exit(0)
+
+    def is_server_running(self):
+        return self.process is not None
+
+    def is_server_restarting(self):
+        return self.restarting
+
+    def launch(self, model: str, launch_kwargs: Optional[Dict[str, Any]] = None):
+        if self.is_server_running():
+            print("Server is already launched.")
+            return
+
+        launch_kwargs = launch_kwargs or self.launch_kwargs
+
+        prefixes = ["openai/", "huggingface/", "local:", "arbor:"]
+        for prefix in prefixes:
+            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 --guided-decoding-backend xgrammar"
+        command = f"python -m sglang_router.launch_server --model-path {model} --dp-size {n_gpus} --router-policy round_robin --port {port} --host 0.0.0.0"
+        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}!")
+
+        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
+
+    def kill(self):
+        from sglang.utils import terminate_process
+
+        if self.process is None:
+            print("No running server to kill.")
+            return
+
+        process = self.process
+        thread = self.thread
+
+        terminate_process(process)
+
+        # Clear references first
+        self.process = None
+        self.thread = None
+        self.get_logs = None
+        self.last_activity = None
+
+        try:
+            # Handle nested signal case
+            if self._shutting_down:
+                process.kill()  # Go straight to SIGKILL if we're shutting down
+            else:
+                process.terminate()  # Try SIGTERM first
+                try:
+                    process.wait(timeout=10)
+                except subprocess.TimeoutExpired:
+                    print(
+                        "Process did not terminate after 10 seconds, forcing with SIGKILL..."
+                    )
+                    process.kill()
+
+            process.wait(timeout=5)
+
+            if thread and thread.is_alive():
+                thread.join(timeout=5)
+
+        except Exception as e:
+            print(f"Error during cleanup: {e}")
+            try:
+                process.kill()  # Final attempt to kill
+            except:
+                pass
+
+        print("Server killed.")
+
+    def run_inference(self, request_json: dict):
+        model = request_json["model"]
+        prefixes = ["openai/", "huggingface/", "local:", "arbor:"]
+        for prefix in prefixes:
+            if model.startswith(prefix):
+                model = model[len(prefix) :]
+        print(f"Running inference for model {model}")
+        # Monkeypatch:
+        if model != self.current_model:
+            print(f"MONKEYPATCH: Model changed from {model} to {self.current_model}")
+            model = self.current_model
+            request_json["model"] = model
+
+        # Update last_activity timestamp
+        self.last_activity = datetime.now()
+
+        if self.process is None or self.launch_kwargs.get("api_base") is None:
+            raise RuntimeError("Server is not running. Please launch it first.")
+
+        if self.restarting:
+            while self.restarting:
+                print("Inference is paused while server is restarting...")
+                time.sleep(5)
+            request_json["model"] = self.current_model
+
+        url = f"{self.launch_kwargs['api_base']}/chat/completions"
+        try:
+            self.inference_count += 1
+            response = requests.post(url, json=request_json)
+            return response.json()
+        except requests.exceptions.ConnectionError:
+            print("Server disconnected...ignoring")
+            return None
+        except Exception as e:
+            print(f"Error during inference: {e}")
+            raise
+        finally:
+            self.inference_count -= 1
+
+    def update_model(self, output_dir):
+        print("Restarting server with new model...")
+        self.restarting = True
+
+        while self.inference_count > 0:
+            print(
+                f"Waiting for inference requests to finish... {self.inference_count} remaining"
+            )
+            time.sleep(5)
+
+        tik = time.time()
+        self.kill()
+        print("Just killed server")
+        # 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):
+            raise RuntimeError(
+                f"Failed to save model - output directory {output_dir} does not exist"
+            )
+
+        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")
+
+
+def get_free_port() -> int:
+    """
+    Return a free TCP port on localhost.
+    """
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("localhost", 0))
+        return s.getsockname()[1]
+
+
+def wait_for_server(base_url: str, timeout: int = None) -> None:
+    """
+    Wait for the server to be ready by polling the /v1/models endpoint.
+
+    Args:
+        base_url: The base URL of the server (e.g. http://localhost:1234)
+        timeout: Maximum time to wait in seconds. None means wait forever.
+    """
+    start_time = time.time()
+    while True:
+        try:
+            response = requests.get(
+                f"{base_url}/v1/models",
+                headers={"Authorization": "Bearer None"},
+            )
+            if response.status_code == 200:
+                # A small extra sleep to ensure server is fully up.
+                time.sleep(5)
+                break
+
+            if timeout and (time.time() - start_time) > timeout:
+                raise TimeoutError("Server did not become ready within timeout period")
+        except requests.exceptions.RequestException:
+            # Server not up yet, wait and retry
+            time.sleep(1)

arbor/server/services/job_manager.py

@@ -0,0 +1,81 @@
+import uuid
+from datetime import datetime
+from typing import Literal
+
+from arbor.server.api.models.schemas import JobStatus
+from arbor.server.core.config import Settings
+
+
+class JobEvent:
+    def __init__(
+        self, level: Literal["info", "warning", "error"], message: str, data: dict = {}
+    ):
+        self.level = level
+        self.message = message
+        self.data = data
+
+        self.id = str(f"ftevent-{uuid.uuid4()}")
+        self.created_at = datetime.now()
+
+
+class JobCheckpoint:
+    def __init__(
+        self,
+        fine_tuned_model_checkpoint: str,
+        fine_tuning_job_id: str,
+        metrics: dict,
+        step_number: int,
+    ):
+        self.id = str(f"ftckpt-{uuid.uuid4()}")
+        self.fine_tuned_model_checkpoint = fine_tuned_model_checkpoint
+        self.fine_tuning_job_id = fine_tuning_job_id
+        self.metrics = metrics
+        self.step_number = step_number
+        self.created_at = datetime.now()
+
+
+class Job:
+    def __init__(self, status: JobStatus):
+        self.id = str(f"ftjob-{uuid.uuid4()}")
+        self.status = status
+        self.fine_tuned_model = None
+        self.events: list[JobEvent] = []
+        self.checkpoints: list[JobCheckpoint] = []
+
+        self.created_at = datetime.now()
+
+    def add_event(self, event: JobEvent):
+        self.events.append(event)
+
+    def get_events(self) -> list[JobEvent]:
+        return self.events
+
+    def add_checkpoint(self, checkpoint: JobCheckpoint):
+        self.checkpoints.append(checkpoint)
+
+    def get_checkpoints(self) -> list[JobCheckpoint]:
+        return self.checkpoints
+
+
+class JobManager:
+    def __init__(self, settings: Settings):
+        self.jobs = {}
+
+    def get_job(self, job_id: str):
+        if job_id not in self.jobs:
+            raise ValueError(f"Job {job_id} not found")
+        return self.jobs[job_id]
+
+    def create_job(self):
+        job = Job(status=JobStatus.PENDING)
+        self.jobs[job.id] = job
+        return job
+
+    def get_jobs(self):
+        return list(self.jobs.values())
+
+    def get_active_job(self):
+        for job in self.jobs.values():
+            if job.status == JobStatus.RUNNING:
+                return job
+        return None