arbor-ai 0.2__tar.gz → 0.2.2__tar.gz

{arbor_ai-0.2 → arbor_ai-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arbor-ai
-Version: 0.2
+Version: 0.2.2
 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
@@ -14,7 +14,7 @@ Requires-Dist: uvicorn
 Requires-Dist: click
 Requires-Dist: python-multipart
 Requires-Dist: pydantic-settings
-Requires-Dist: vllm>=0.8.5.post1
+Requires-Dist: vllm==0.8.5.post1
 Requires-Dist: transformers
 Requires-Dist: trl>=0.17.0
 Requires-Dist: peft
@@ -43,7 +43,8 @@ Install Arbor via pip:
 pip install -U arbor-ai
 ```
 
-Optionally, you can also install:
+Optionally, you can also install flash attention to speed up inference. <br/>
+This can take 15+ minutes to install on some setups:
 ```bash
 pip install flash-attn --no-build-isolation
 ```
@@ -52,33 +53,32 @@ pip install flash-attn --no-build-isolation
 
 ## ⚡ Quick Start
 
-### 1️⃣ Make an `arbor.yaml` File
-
-This is all dependent on your setup. Here is an example of one:
-```yaml
-inference:
-  gpu_ids: '0'
-
-training:
-  gpu_ids: '1, 2'
-```
-Which will use the `GPU:0` for inference with `GPU:1` and `GPU:2` reserved for training. We generally recommend splitting the GPUs roughly evenly between inference and training.
-
-### 2️⃣ Start the Server
+### 1️⃣ Start the Server
 
 **CLI:**
 
 ```bash
-python -m arbor.cli serve --arbor-config arbor.yaml
+python -m arbor.cli serve
 ```
+On the first run you'll be asked which GPUs will be used for training and which for inference. For more that one GPU, separate the ids by comma: `1, 2`. Your config file will be saved in `~/.arbor/config.yaml` should you want to edit these configs in the future.
 
-### 3️⃣ Optimize a DSPy Program
+### 2️⃣ Optimize a DSPy Program
 
 Follow the DSPy tutorials here to see usage examples:
 [DSPy RL Optimization Examples](https://dspy.ai/tutorials/rl_papillon/)
 
 ---
 
+### Troubleshooting
+
+**NCCL Errors**
+Certain GPU setups, particularly with newer GPUs, seem to have issues with NCCL that cause Arbor to crash. Often times of these can be fixed with the following environment variables:
+
+```bash
+export NCCL_P2P_DISABLE=1
+export NCCL_IB_DISABLE=1
+```
+
 ## 🙏 Acknowledgements
 
 Arbor builds on the shoulders of great work. We extend our thanks to:

{arbor_ai-0.2 → arbor_ai-0.2.2}/README.md

@@ -16,7 +16,8 @@ Install Arbor via pip:
 pip install -U arbor-ai
 ```
 
-Optionally, you can also install:
+Optionally, you can also install flash attention to speed up inference. <br/>
+This can take 15+ minutes to install on some setups:
 ```bash
 pip install flash-attn --no-build-isolation
 ```
@@ -25,33 +26,32 @@ pip install flash-attn --no-build-isolation
 
 ## ⚡ Quick Start
 
-### 1️⃣ Make an `arbor.yaml` File
-
-This is all dependent on your setup. Here is an example of one:
-```yaml
-inference:
-  gpu_ids: '0'
-
-training:
-  gpu_ids: '1, 2'
-```
-Which will use the `GPU:0` for inference with `GPU:1` and `GPU:2` reserved for training. We generally recommend splitting the GPUs roughly evenly between inference and training.
-
-### 2️⃣ Start the Server
+### 1️⃣ Start the Server
 
 **CLI:**
 
 ```bash
-python -m arbor.cli serve --arbor-config arbor.yaml
+python -m arbor.cli serve
 ```
+On the first run you'll be asked which GPUs will be used for training and which for inference. For more that one GPU, separate the ids by comma: `1, 2`. Your config file will be saved in `~/.arbor/config.yaml` should you want to edit these configs in the future.
 
-### 3️⃣ Optimize a DSPy Program
+### 2️⃣ Optimize a DSPy Program
 
 Follow the DSPy tutorials here to see usage examples:
 [DSPy RL Optimization Examples](https://dspy.ai/tutorials/rl_papillon/)
 
 ---
 
+### Troubleshooting
+
+**NCCL Errors**
+Certain GPU setups, particularly with newer GPUs, seem to have issues with NCCL that cause Arbor to crash. Often times of these can be fixed with the following environment variables:
+
+```bash
+export NCCL_P2P_DISABLE=1
+export NCCL_IB_DISABLE=1
+```
+
 ## 🙏 Acknowledgements
 
 Arbor builds on the shoulders of great work. We extend our thanks to:

arbor_ai-0.2.2/arbor/__init__.py

@@ -0,0 +1,17 @@
+"""
+Arbor - A framework for fine-tuning and managing language models
+"""
+
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("arbor-ai")
+except PackageNotFoundError:
+    # Package is not installed, likely in development mode
+    __version__ = "dev"
+except Exception:
+    __version__ = "unknown"
+
+from arbor.client.arbor_client import is_running, serve, stop
+
+__all__ = ["__version__", "serve", "stop", "is_running"]

arbor_ai-0.2.2/arbor/cli.py

@@ -0,0 +1,153 @@
+import os
+from datetime import datetime
+
+import click
+import uvicorn
+
+from arbor.server.core.config import Config
+from arbor.server.core.config_manager import ConfigManager
+from arbor.server.main import app
+from arbor.server.services.file_manager import FileManager
+from arbor.server.services.grpo_manager import GRPOManager
+from arbor.server.services.health_manager import HealthManager
+from arbor.server.services.inference_manager import InferenceManager
+from arbor.server.services.job_manager import JobManager
+from arbor.server.services.training_manager import TrainingManager
+from arbor.client.arbor_client import create_app
+
+
+@click.group()
+def cli():
+    pass
+
+def start_server(host="0.0.0.0", port=7453, storage_path="./storage", timeout=10):
+    """Start the Arbor API server with a single function call"""
+    import socket
+    import threading
+    import time
+    from contextlib import closing
+
+    def is_port_in_use(port):
+        with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
+            return sock.connect_ex(("localhost", port)) == 0
+
+    # First ensure the port is free
+    if is_port_in_use(port):
+        raise RuntimeError(f"Port {port} is already in use")
+
+    app = create_app(storage_path)
+    # configure_uvicorn_logging()
+    config = uvicorn.Config(app, host=host, port=port, log_level="info")
+    server = uvicorn.Server(config)
+
+    def run_server():
+        server.run()
+
+    thread = threading.Thread(target=run_server, daemon=True)
+    thread.start()
+
+    # Wait for server to start
+    start_time = time.time()
+    while not is_port_in_use(port):
+        if time.time() - start_time > timeout:
+            raise TimeoutError(f"Server failed to start within {timeout} seconds")
+        time.sleep(0.1)
+
+    # Give it a little extra time to fully initialize
+    time.sleep(0.5)
+
+    return server
+
+
+def stop_server(server):
+    """Stop the Arbor API server"""
+    server.should_exit = True
+
+
+@cli.command()
+@click.option("--host", default="0.0.0.0", help="Host to bind to")
+@click.option("--port", default=7453, help="Port to bind to")
+@click.option("--arbor-config", required=False, help="Path to the Arbor config file")
+def serve(host, port, arbor_config):
+    """Start the Arbor API server"""
+
+    if arbor_config:
+        config_path = arbor_config
+    else:
+        config_path = Config.use_default_config()
+
+        # If no config found, run first-time setup
+        if config_path is None:
+            config_path = run_first_time_setup()
+
+    # Validate config exists and is readable
+    is_valid, msg = ConfigManager.validate_config_file(config_path)
+
+    if not is_valid:
+        click.echo(msg)
+        raise click.Abort()
+
+    try:
+        create_app(config_path)
+        # Temporarily disable custom uvicorn logging configuration
+        # configure_uvicorn_logging()
+        uvicorn.run(app, host=host, port=port)
+    except Exception as e:
+        click.echo(f"Failed to start server: {e}", err=True)
+        raise click.Abort()
+
+
+def run_first_time_setup() -> str:
+    """Run first-time setup and return created config path"""
+    click.echo("Welcome to Arbor!")
+    click.echo("It looks like this is your first time running Arbor.")
+    click.echo("Let's set up your configuration...\n")
+
+    try:
+        # Get config details
+        inference = click.prompt(
+            "Which gpu ids should be used for inference (separated by comma)",
+            default="0",
+        )
+        training = click.prompt(
+            "Which gpu ids should be used for training (separated by comma)",
+            default="1, 2",
+        )
+        click.echo()
+
+        # Get config file path
+        config_path = click.prompt(
+            "Enter path to save config file in. We recommend (~/.arbor/config.yaml)",
+            default=ConfigManager.get_default_config_path(),
+        )
+        logger = get_logger(__name__)
+        logger.info(f"Config path selected: {config_path}")
+        click.echo()
+
+        # Update or create config at path
+        config_path = ConfigManager.update_config(inference, training, config_path)
+        click.echo(f"Created configuration at: {config_path}")
+
+        # Check if it is a valid config file
+        is_valid, msg = ConfigManager.validate_config_file(config_path)
+        if not is_valid:
+            raise click.ClickException(f"Invalid config file: {msg}")
+
+        # Read and display the contents
+        _, content = ConfigManager.get_config_contents(config_path)
+
+        click.echo("\nConfiguration file contents:")
+        click.echo("---")
+        click.echo(content)
+        click.echo("---")
+
+        click.echo("\nSetup complete! Starting Arbor server...")
+        return config_path
+
+    except Exception as e:
+        click.echo(f"Failed initial setup of Arbor: {e}", err=True)
+        raise click.Abort()
+
+
+if __name__ == "__main__":
+    cli()

arbor_ai-0.2.2/arbor/client/arbor_client.py

@@ -0,0 +1,259 @@
+import asyncio
+import os
+import socket
+import threading
+import time
+from datetime import datetime
+
+import click
+import requests
+import uvicorn
+
+from arbor.server.core.config import Config
+from arbor.server.core.config_manager import ConfigManager
+from arbor.server.main import app
+from arbor.server.services.file_manager import FileManager
+from arbor.server.services.grpo_manager import GRPOManager
+from arbor.server.services.health_manager import HealthManager
+from arbor.server.services.inference_manager import InferenceManager
+from arbor.server.services.job_manager import JobManager
+from arbor.server.services.training_manager import TrainingManager
+
+# Global server state
+_server = None
+_server_thread = None
+_server_loop = None
+_server_host = None
+_server_port = None
+
+
+def create_app(
+    config_path: str = None,
+    storage_path: str = None,
+    inference_gpus: str = None,
+    training_gpus: str = None,
+):
+    """Create and configure the Arbor API application
+
+    Args:
+        arbor_config_path (str): Path to config file
+        storage_path (str): Path to storage directory
+        inference_gpus (str): gpu ids to use for inference
+        training_gpus (str): gpu ids to use for training
+
+    Returns:
+        FastAPI: Configured FastAPI application
+    """
+    # Create new config instance with overrides
+    if config_path:
+        config = Config.load_config_from_yaml(config_path)
+    elif inference_gpus and training_gpus:
+        config = Config.load_config_directly(
+            storage_path, inference_gpus, training_gpus
+        )
+    else:
+        raise ValueError(
+            "Either 'config_path' must be provided, or 'inference_gpus', and 'training_gpus' must be provided"
+        )
+
+    app.state.log_dir = Config.make_log_dir(config.STORAGE_PATH)
+
+    # Initialize services with config
+    health_manager = HealthManager(config=config)
+    file_manager = FileManager(config=config)
+    job_manager = JobManager(config=config)
+    training_manager = TrainingManager(config=config)
+    inference_manager = InferenceManager(config=config)
+    grpo_manager = GRPOManager(config=config)
+
+    # Inject config into app state
+    app.state.config = config
+    app.state.file_manager = file_manager
+    app.state.job_manager = job_manager
+    app.state.training_manager = training_manager
+    app.state.inference_manager = inference_manager
+    app.state.grpo_manager = grpo_manager
+    app.state.health_manager = health_manager
+
+    return app
+
+
+def serve(
+    config_path: str = None,
+    storage_path: str = None,
+    inference_gpus: str = None,
+    training_gpus: str = None,
+    host: str = "0.0.0.0",
+    port: int = 7453,
+):
+    """Start the Arbor API server.
+
+    Starts the server in a background thread and returns once the server is ready to accept requests.
+    Use arbor.stop() to shutdown the server.
+
+    Args:
+        config_path: Path to YAML config file (optional)
+        storage_path: Valid storage directory path (optional)
+        inference_gpus: GPU IDs for inference, e.g. "0,1" (optional, default 0)
+        training_gpus: GPU IDs for training, e.g. "1,2,3" (optional, default 1,2)
+        host: Host to bind to (default: "0.0.0.0")
+        port: Port to bind to (default: 7453)
+
+    Example:
+        import arbor
+        arbor.serve(inference_gpus="0", training_gpus="1,2")
+        # Server is now ready to accept requests
+        # Later, to stop:
+        arbor.stop()
+    """
+    global _server, _server_thread, _server_loop, _server_host, _server_port
+
+    # Stop existing server if running
+    if _server is not None:
+        print("🌳 Stopping existing server...")
+        stop()
+
+    _server_host = host
+    _server_port = port
+
+    create_app(config_path, storage_path, inference_gpus, training_gpus)
+
+    # Start server in background thread
+    def run_server():
+        global _server, _server_loop
+
+        # Create a new event loop for this thread
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        _server_loop = loop
+
+        # Create uvicorn config and server
+        config = uvicorn.Config(app, host=host, port=port, loop=loop)
+        server = uvicorn.Server(config)
+        _server = server
+
+        # Run the server
+        try:
+            loop.run_until_complete(server.serve())
+        except Exception as e:
+            print(f"Server error: {e}")
+        finally:
+            loop.close()
+            _server = None
+            _server_loop = None
+
+    # Start server thread
+    _server_thread = threading.Thread(target=run_server, daemon=True)
+    _server_thread.start()
+
+    print(f"🌳 Starting Arbor server on http://{host}:{port}...")
+
+    # Wait for server to be ready
+    try:
+        _wait_for_server_ready(host, port, timeout=60)  # Increased timeout
+        print(f"🌳 Arbor server is ready and accepting requests!")
+    except TimeoutError as e:
+        print(f"❌ {e}")
+        # Try to stop the server if it failed to start properly
+        stop()
+        raise
+
+
+def stop():
+    """Stop the Arbor server if it's running."""
+    global _server, _server_thread, _server_loop
+
+    if _server is None:
+        print("🌳 No server running to stop.")
+        return
+
+    print("🌳 Stopping Arbor server...")
+
+    # Schedule server shutdown in the server's event loop
+    if _server_loop and _server:
+        try:
+            asyncio.run_coroutine_threadsafe(_server.shutdown(), _server_loop)
+        except Exception as e:
+            print(f"Error during shutdown: {e}")
+
+    # Wait for thread to finish
+    if _server_thread and _server_thread.is_alive():
+        _server_thread.join(timeout=5)
+
+    # Reset global state
+    _server = None
+    _server_thread = None
+    _server_loop = None
+
+    print("🌳 Arbor server stopped.")
+
+
+def is_running():
+    """Check if the Arbor server is currently running."""
+    return (
+        _server is not None and _server_thread is not None and _server_thread.is_alive()
+    )
+
+
+def _wait_for_server_ready(host, port, timeout=30):
+    """Wait for the server to be ready to accept requests."""
+    start_time = time.time()
+    last_error = None
+    port_open = False
+
+    print(f"🌳 Waiting for server to be ready at http://{host}:{port}...")
+
+    while time.time() - start_time < timeout:
+        # First check if the port is open
+        if not port_open:
+            try:
+                sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+                sock.settimeout(1)
+                result = sock.connect_ex((host, port))
+                sock.close()
+                if result == 0:
+                    port_open = True
+                    print(f"🌳 Port {port} is now open, checking health endpoint...")
+                else:
+                    last_error = f"Port {port} not yet open"
+                    time.sleep(0.5)
+                    continue
+            except Exception as e:
+                last_error = f"Socket error: {e}"
+                time.sleep(0.5)
+                continue
+
+        # Now try the health check
+        try:
+            response = requests.get(f"http://{host}:{port}/health/simple", timeout=2)
+            if response.status_code == 200:
+                print(f"🌳 Server ready! Response: {response.json()}")
+                return
+            else:
+                last_error = f"Health check returned status {response.status_code}"
+        except requests.exceptions.ConnectionError as e:
+            last_error = f"Connection error: {e}"
+            port_open = False  # Port might have closed
+        except requests.exceptions.Timeout as e:
+            last_error = f"Timeout error: {e}"
+        except requests.exceptions.RequestException as e:
+            last_error = f"Request error: {e}"
+        except Exception as e:
+            last_error = f"Unexpected error: {e}"
+
+        # Print progress every 5 seconds
+        elapsed = time.time() - start_time
+        if int(elapsed) % 5 == 0 and elapsed >= 5:
+            print(
+                f"🌳 Still waiting... ({elapsed:.1f}s elapsed, port_open={port_open}, last error: {last_error})"
+            )
+
+        time.sleep(0.5)
+
+    raise TimeoutError(
+        f"Server did not become ready within {timeout} seconds. Last error: {last_error}"
+    )
+
+
+if __name__ == "__main__":
+    serve(inference_gpus="0, 1", training_gpus="2, 3")

{arbor_ai-0.2 → arbor_ai-0.2.2}/arbor/server/api/models/schemas.py

@@ -178,7 +178,7 @@ class ChatCompletionModel(BaseModel):
 
 class GRPORequest(BaseModel):
     model: str
-    batch: List[dict]
+    batch: List[dict] | List[List[dict]]
 
 
 class GRPOConfigRequest(BaseModel):
@@ -205,6 +205,8 @@ class GRPOConfigRequest(BaseModel):
     # Arbor specific
     max_context_length: Optional[int] = None
     lora: Optional[bool] = None
+    grpo_flavor: Optional[Literal["grpo", "mmgrpo"]] = None
+    wandb_kwargs: Optional[dict] = None
     # To name the run
     suffix: Optional[str] = None
     generation_batch_size: Optional[int] = None

{arbor_ai-0.2 → arbor_ai-0.2.2}/arbor/server/api/routes/grpo.py

@@ -10,7 +10,6 @@ from arbor.server.api.models.schemas import (
     GRPOConfigResponse,
     GRPORequest,
     GRPOStepResponse,
-    GRPOTerminateRequest,
     GRPOTerminateResponse,
 )
 
@@ -27,12 +26,9 @@ def initialize_grpo(request: Request, grpo_config_request: GRPOConfigRequest):
 
 # Create a grpo job
 @router.post("/step", response_model=GRPOStepResponse)
-def run_grpo_step(
-    request: Request, grpo_request: GRPORequest, background_tasks: BackgroundTasks
-):
-    inference_manager = request.app.state.inference_manager
+def run_grpo_step(request: Request, grpo_request: GRPORequest):
     grpo_manager = request.app.state.grpo_manager
-
+    inference_manager = request.app.state.inference_manager
     step_data = grpo_manager.grpo_step(grpo_request, inference_manager)
 
     return GRPOStepResponse(status="success", **step_data)

{arbor_ai-0.2 → arbor_ai-0.2.2}/arbor/server/api/routes/inference.py

@@ -3,6 +3,10 @@ import uuid
 
 from fastapi import APIRouter, Request
 
+from arbor.server.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
 router = APIRouter()
 
 
@@ -27,17 +31,17 @@ async def run_inference(
 
     # if a server isnt running, launch one
     if not inference_manager.is_server_running():
-        print("No model is running, launching model...")
+        logger.info("No model is running, launching model...")
         inference_manager.launch(request_model)
 
     # if the requested model is different from the launched model, swap the server
     if request_model != inference_manager.launched_model:
-        print(
+        logger.info(
             f"Model changed from {inference_manager.launched_model} to {request_model}, swapping server..."
         )
         inference_manager.kill()
         inference_manager.launch(request_model)
-        print(f"Model swapped to {request_model}")
+        logger.info(f"Model swapped to {request_model}")
 
     # forward the request to the inference server
     completion = await inference_manager.run_inference(raw_json)