arbor-ai 0.1.4__tar.gz → 0.1.6__tar.gz

arbor_ai-0.1.6/PKG-INFO

@@ -0,0 +1,78 @@
+Metadata-Version: 2.4
+Name: arbor-ai
+Version: 0.1.6
+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
+Project-URL: Issues, https://github.com/Ziems/arbor/issues
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: fastapi
+Requires-Dist: uvicorn
+Requires-Dist: click
+Requires-Dist: python-multipart
+Requires-Dist: pydantic-settings
+Requires-Dist: torch
+Requires-Dist: transformers
+Requires-Dist: trl
+Requires-Dist: peft
+Requires-Dist: ray>=2.9
+Requires-Dist: setuptools<77.0.0,>=76.0.0
+Requires-Dist: pyzmq>=26.4.0
+Requires-Dist: pyyaml>=6.0.2
+Requires-Dist: sglang>=0.4.5.post3
+Requires-Dist: sglang-router
+Dynamic: license-file
+
+<p align="center">
+  <img src="https://github.com/user-attachments/assets/ed0dd782-65fa-48b5-a762-b343b183be09" alt="Description" width="400"/>
+</p>
+
+**A framework for optimizing DSPy programs with RL.**
+
+---
+
+## 🚀 Installation
+
+Install Arbor via pip:
+
+```bash
+pip install git+https://github.com/Ziems/arbor.git
+```
+
+---
+
+## ⚡ 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'
+```
+
+### 2️⃣ Start the Server
+
+**CLI:**
+
+```bash
+python -m arbor.cli serve --arbor-config arbor.yaml
+```
+
+### 3️⃣ Optimize a DSPy Program
+
+Follow the DSPy tutorials here to see usage examples:
+[DSPy RL Optimization Examples](https://dspy.ai/tutorials/rl_papillon/)
+
+---
+
+## 🙏 Acknowledgements
+
+Arbor builds on the shoulders of great work. We extend our thanks to:
+- **[Will Brown's Verifiers library](https://github.com/willccbb/verifiers)**
+- **[Hugging Face TRL library](https://github.com/huggingface/trl)**

arbor_ai-0.1.6/README.md

@@ -0,0 +1,51 @@
+<p align="center">
+  <img src="https://github.com/user-attachments/assets/ed0dd782-65fa-48b5-a762-b343b183be09" alt="Description" width="400"/>
+</p>
+
+**A framework for optimizing DSPy programs with RL.**
+
+---
+
+## 🚀 Installation
+
+Install Arbor via pip:
+
+```bash
+pip install git+https://github.com/Ziems/arbor.git
+```
+
+---
+
+## ⚡ 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'
+```
+
+### 2️⃣ Start the Server
+
+**CLI:**
+
+```bash
+python -m arbor.cli serve --arbor-config arbor.yaml
+```
+
+### 3️⃣ Optimize a DSPy Program
+
+Follow the DSPy tutorials here to see usage examples:
+[DSPy RL Optimization Examples](https://dspy.ai/tutorials/rl_papillon/)
+
+---
+
+## 🙏 Acknowledgements
+
+Arbor builds on the shoulders of great work. We extend our thanks to:
+- **[Will Brown's Verifiers library](https://github.com/willccbb/verifiers)**
+- **[Hugging Face TRL library](https://github.com/huggingface/trl)**

arbor_ai-0.1.6/arbor/cli.py

@@ -0,0 +1,101 @@
+import click
+import uvicorn
+
+from arbor.server.core.config import Settings
+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.inference_manager import InferenceManager
+from arbor.server.services.job_manager import JobManager
+from arbor.server.services.training_manager import TrainingManager
+
+
+@click.group()
+def cli():
+    pass
+
+
+def create_app(arbor_config_path: str):
+    """Create and configure the Arbor API application
+
+    Args:
+        storage_path (str): Path to store models and uploaded training files
+
+    Returns:
+        FastAPI: Configured FastAPI application
+    """
+    # Create new settings instance with overrides
+    settings = Settings.load_from_yaml(arbor_config_path)
+
+    # Initialize services with settings
+    file_manager = FileManager(settings=settings)
+    job_manager = JobManager(settings=settings)
+    training_manager = TrainingManager(settings=settings)
+    inference_manager = InferenceManager(settings=settings)
+    grpo_manager = GRPOManager(settings=settings)
+    # Inject settings into app state
+    app.state.settings = settings
+    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
+
+    return app
+
+
+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)
+    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=True, help="Path to the Arbor config file")
+def serve(host, port, arbor_config):
+    """Start the Arbor API server"""
+    app = create_app(arbor_config)
+    uvicorn.run(app, host=host, port=port)
+
+
+if __name__ == "__main__":
+    cli()

arbor_ai-0.1.6/arbor/client/api.py

@@ -0,0 +1 @@
+# Unused Right Now

arbor_ai-0.1.6/arbor/server/api/models/schemas.py

@@ -0,0 +1,223 @@
+from enum import Enum
+from typing import Any, Generic, List, Literal, Optional, TypeVar
+
+from pydantic import BaseModel, ConfigDict
+
+# Generic type for list items
+T = TypeVar("T")
+
+
+class PaginatedResponse(BaseModel, Generic[T]):
+    object: str = "list"
+    data: List[T]
+    has_more: bool = False
+
+
+class FileModel(BaseModel):
+    id: str
+    object: str = "file"
+    bytes: int
+    created_at: int
+    filename: str
+    purpose: str
+
+
+class WandbConfig(BaseModel):
+    project: str
+    name: Optional[str] = None
+    entity: Optional[str] = None
+    tags: Optional[List[str]] = None
+
+
+class IntegrationModel(BaseModel):
+    type: str
+    wandb: WandbConfig
+
+
+class FineTuneRequest(BaseModel):
+    model: str
+    training_file: str  # id of uploaded jsonl file
+    method: dict
+    suffix: Optional[str] = None
+    # UNUSED
+    validation_file: Optional[str] = None
+    integrations: Optional[List[IntegrationModel]] = []
+    seed: Optional[int] = None
+
+
+class ErrorModel(BaseModel):
+    code: str
+    message: str
+    param: str | None = None
+
+
+class SupervisedHyperparametersModel(BaseModel):
+    batch_size: int | str = "auto"
+    learning_rate_multiplier: float | str = "auto"
+    n_epochs: int | str = "auto"
+
+
+class DPOHyperparametersModel(BaseModel):
+    beta: float | str = "auto"
+    batch_size: int | str = "auto"
+    learning_rate_multiplier: float | str = "auto"
+    n_epochs: int | str = "auto"
+
+
+class SupervisedModel(BaseModel):
+    hyperparameters: SupervisedHyperparametersModel
+
+
+class DpoModel(BaseModel):
+    hyperparameters: DPOHyperparametersModel
+
+
+class MethodModel(BaseModel):
+    type: Literal["supervised"] | Literal["dpo"]
+    supervised: SupervisedModel | None = None
+    dpo: DpoModel | None = None
+
+
+# https://platform.openai.com/docs/api-reference/fine-tuning/object
+class JobStatus(Enum):
+    PENDING = "pending"  # Not in OAI
+    PENDING_PAUSE = "pending_pause"  # Not in OAI
+    PENDING_RESUME = "pending_resume"  # Not in OAI
+    PAUSED = "paused"  # Not in OAI
+    VALIDATING_FILES = "validating_files"
+    QUEUED = "queued"
+    RUNNING = "running"
+    SUCCEEDED = "succeeded"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+    PENDING_CANCEL = "pending_cancel"
+
+
+# https://platform.openai.com/docs/api-reference/fine-tuning/object
+class JobStatusModel(BaseModel):
+    object: str = "fine_tuning.job"
+    id: str
+    fine_tuned_model: str | None = None
+    status: JobStatus
+
+    # UNUSED so commented out
+    # model: str
+    # created_at: int
+    # error: ErrorModel | None = None
+    # details: str = ""
+    # finished_at: int
+    # hyperparameters: None # deprecated in OAI
+    # organization_id: str
+    # result_files: list[str]
+    # trained_tokens: int | None = None # None if not finished
+    # training_file: str
+    # validation_file: str
+    # integrations: list[Integration]
+    # seed: int
+    # estimated_finish: int | None = None # The Unix timestamp (in seconds) for when the fine-tuning job is estimated to finish. The value will be null if the fine-tuning job is not running.
+    # method: MethodModel
+    # metadata: dict[str, str]
+
+
+class JobEventModel(BaseModel):
+    object: str = "fine_tuning.job_event"
+    id: str
+    created_at: int
+    level: str
+    message: str
+    data: dict[str, Any]
+    type: str
+
+
+class MetricsModel(BaseModel):
+    step: int
+    train_loss: float
+    train_mean_token_accuracy: float
+    valid_loss: float
+    valid_mean_token_accuracy: float
+    full_valid_loss: float
+    full_valid_mean_token_accuracy: float
+
+
+class JobCheckpointModel(BaseModel):
+    object: str = "fine_tuning.job_checkpoint"
+    id: str
+    created_at: int
+    fine_tuned_model_checkpoint: str
+    step_number: int
+    metrics: MetricsModel
+    fine_tuning_job_id: str
+
+
+class ChatCompletionMessage(BaseModel):
+    role: Literal["system", "user", "assistant"]
+    content: str
+
+
+class ChatCompletionRequest(BaseModel):
+    model: str
+    messages: List[ChatCompletionMessage]
+    temperature: float | None = None
+    top_p: float | None = None
+    max_tokens: int | None = None
+
+
+class ChatCompletionChoice(BaseModel):
+    message: ChatCompletionMessage
+    index: int
+    finish_reason: Literal["stop", "length", "tool_calls"]
+
+
+class ChatCompletionModel(BaseModel):
+    id: str
+    object: str = "chat.completion"
+    created: int
+    model: str
+    choices: List[ChatCompletionChoice]
+
+
+class GRPORequest(BaseModel):
+    model: str
+    update_inference_model: bool
+    batch: List[dict]
+
+
+class GRPOConfigRequest(BaseModel):
+    model: str
+    temperature: Optional[float] = None
+    beta: Optional[float] = None
+    num_iterations: Optional[int] = None
+    num_generations: Optional[int] = None
+    per_device_train_batch_size: Optional[int] = None
+    learning_rate: Optional[float] = None
+    gradient_accumulation_steps: Optional[int] = None
+    gradient_checkpointing: Optional[bool] = None
+    lr_scheduler_type: Optional[str] = None
+    max_prompt_length: Optional[int] = None
+    max_completion_length: Optional[int] = None
+    gradient_checkpointing_kwargs: Optional[dict] = {}
+    bf16: Optional[bool] = None
+    scale_rewards: Optional[bool] = None
+    max_grad_norm: Optional[float] = None
+    lora: Optional[bool] = None
+    update_interval: Optional[int] = None
+    # To name the run
+    suffix: Optional[str] = None
+
+
+class GRPOConfigResponse(BaseModel):
+    status: str
+
+
+class GRPOTerminateRequest(BaseModel):
+    status: Optional[str] = "success"
+
+
+class GRPOTerminateResponse(BaseModel):
+    status: str
+    current_model: str
+
+
+class GRPOStepResponse(BaseModel):
+    status: str
+    current_model: str

arbor_ai-0.1.6/arbor/server/api/routes/files.py

@@ -0,0 +1,52 @@
+from typing import Literal
+
+from fastapi import APIRouter, Body, File, HTTPException, Request, UploadFile
+
+from arbor.server.api.models.schemas import FileModel, PaginatedResponse
+from arbor.server.services.file_manager import FileValidationError
+
+# https://platform.openai.com/docs/api-reference/files/list
+router = APIRouter()
+
+
+@router.post("", response_model=FileModel)
+async def upload_file(
+    request: Request,
+    file: UploadFile = File(...),
+    purpose: Literal["assistants", "vision", "fine-tune", "batch"] = Body("fine-tune"),
+):
+    file_manager = request.app.state.file_manager
+    if not file.filename.endswith(".jsonl"):
+        raise HTTPException(status_code=400, detail="Only .jsonl files are allowed")
+
+    try:
+        content = await file.read()
+        # file_manager.validate_file_format(content)   #TODO: add another flag to specify the types of files
+        await file.seek(0)  # Reset file pointer to beginning
+        return FileModel(**file_manager.save_uploaded_file(file))
+    except FileValidationError as e:
+        raise HTTPException(status_code=400, detail=f"Invalid file format: {str(e)}")
+
+
+@router.get("", response_model=PaginatedResponse[FileModel])
+def list_files(request: Request):
+    file_manager = request.app.state.file_manager
+    return PaginatedResponse(
+        items=file_manager.get_files(),
+        total=len(file_manager.get_files()),
+        page=1,
+        page_size=10,
+    )
+
+
+@router.get("/{file_id}", response_model=FileModel)
+def get_file(request: Request, file_id: str):
+    file_manager = request.app.state.file_manager
+    return file_manager.get_file(file_id)
+
+
+@router.delete("/{file_id}")
+def delete_file(request: Request, file_id: str):
+    file_manager = request.app.state.file_manager
+    file_manager.delete_file(file_id)
+    return {"message": "File deleted"}

arbor_ai-0.1.6/arbor/server/api/routes/grpo.py

@@ -0,0 +1,54 @@
+import os
+import subprocess
+
+from fastapi import APIRouter, BackgroundTasks, Request
+
+from arbor.server.api.models.schemas import (
+    GRPOConfigRequest,
+    GRPOConfigResponse,
+    GRPORequest,
+    GRPOStepResponse,
+    GRPOTerminateRequest,
+    GRPOTerminateResponse,
+)
+
+router = APIRouter()
+
+
+@router.post("/initialize", response_model=GRPOConfigResponse)
+def initialize_grpo(request: Request, grpo_config_request: GRPOConfigRequest):
+    inference_manager = request.app.state.inference_manager
+    grpo_manager = request.app.state.grpo_manager
+    grpo_manager.initialize(grpo_config_request, inference_manager)
+    return GRPOConfigResponse(status="success")
+
+
+# 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
+    grpo_manager = request.app.state.grpo_manager
+
+    current_model = grpo_manager.grpo_step(grpo_request, inference_manager)
+
+    return GRPOStepResponse(status="success", current_model=current_model)
+
+
+@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)
+
+
+@router.post("/terminate", response_model=GRPOTerminateResponse)
+def terminate_grpo(request: Request):
+    # No body needed for this request at this moment
+    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)

arbor_ai-0.1.6/arbor/server/api/routes/inference.py

@@ -0,0 +1,53 @@
+import time
+
+from fastapi import APIRouter, Request
+
+router = APIRouter()
+
+
+@router.post("/completions")
+async def run_inference(
+    request: Request,
+):
+    inference_manager = request.app.state.inference_manager
+    raw_json = await request.json()
+
+    prefixes = ["openai/", "huggingface/", "local:", "arbor:"]
+    for prefix in prefixes:
+        if raw_json["model"].startswith(prefix):
+            raw_json["model"] = raw_json["model"][len(prefix) :]
+
+    # if a server isnt running, launch one
+    if (
+        not inference_manager.is_server_running()
+        and not inference_manager.is_server_restarting()
+    ):
+        print("No model is running, launching model...")
+        inference_manager.launch(raw_json["model"])
+
+    if inference_manager.is_server_restarting():
+        print("Waiting for server to finish restarting...")
+        while inference_manager.is_server_restarting():
+            time.sleep(5)
+        # Update the model in the request
+        raw_json["model"] = inference_manager.current_model
+
+    # forward the request to the inference server
+    completion = inference_manager.run_inference(raw_json)
+
+    return completion
+
+
+@router.post("/launch")
+async def launch_inference(request: Request):
+    inference_manager = request.app.state.inference_manager
+    raw_json = await request.json()
+    inference_manager.launch(raw_json["model"], raw_json["launch_kwargs"])
+    return {"message": "Inference server launched"}
+
+
+@router.post("/kill")
+async def kill_inference(request: Request):
+    inference_manager = request.app.state.inference_manager
+    inference_manager.kill()
+    return {"message": "Inference server killed"}