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

arbor/cli.py

@@ -1,17 +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=8000, help='Port to bind to')
-def serve(host, port):
+@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()
+
+if __name__ == "__main__":
+    cli()

arbor/client/api.py

@@ -1,2 +1 @@
-from typing import Optional, Dict, Any
-
+# Unused Right Now

arbor/server/api/models/schemas.py

@@ -1,6 +1,19 @@
-from pydantic import BaseModel
+from enum import Enum
+from typing import Any, Generic, List, Literal, Optional, TypeVar
 
-class FileResponse(BaseModel):
+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
@@ -8,12 +21,203 @@ class FileResponse(BaseModel):
     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 JobStatusResponse(BaseModel):
+
+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
-    details: str = ""
-    fine_tuned_model: str | None = None
+    current_model: str

arbor/server/api/routes/files.py

@@ -1,23 +1,52 @@
-from fastapi import APIRouter, UploadFile, File, Depends, HTTPException
-from arbor.server.services.file_manager import FileManager
-from arbor.server.api.models.schemas import FileResponse
-from arbor.server.services.dependencies import get_file_manager
+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=FileResponse)
+
+@router.post("", response_model=FileModel)
 async def upload_file(
+    request: Request,
     file: UploadFile = File(...),
-    file_manager: FileManager = Depends(get_file_manager)
+    purpose: Literal["assistants", "vision", "fine-tune", "batch"] = Body("fine-tune"),
 ):
-    if not file.filename.endswith('.jsonl'):
+    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)
+        # 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 file_manager.save_uploaded_file(file)
+        return FileModel(**file_manager.save_uploaded_file(file))
     except FileValidationError as e:
-        raise HTTPException(status_code=400, detail=f"Invalid file format: {str(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/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/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"}

arbor/server/api/routes/jobs.py

@@ -1,14 +1,117 @@
-from fastapi import APIRouter, Depends
-from arbor.server.services.job_manager import JobManager
-from arbor.server.services.dependencies import get_job_manager
-from arbor.server.api.models.schemas import JobStatusResponse
+from fastapi import APIRouter, BackgroundTasks, HTTPException, Request
+
+from arbor.server.api.models.schemas import (
+    FineTuneRequest,
+    JobCheckpointModel,
+    JobEventModel,
+    JobStatus,
+    JobStatusModel,
+    PaginatedResponse,
+)
+from arbor.server.services.job_manager import JobStatus
 
 router = APIRouter()
 
-@router.get("/{job_id}", response_model=JobStatusResponse)
+
+# Create a fine-tune job
+@router.post("", response_model=JobStatusModel)
+def create_fine_tune_job(
+    request: Request,
+    fine_tune_request: FineTuneRequest,
+    background_tasks: BackgroundTasks,
+):
+    job_manager = request.app.state.job_manager
+    file_manager = request.app.state.file_manager
+    training_manager = request.app.state.training_manager
+
+    job = job_manager.create_job()
+    background_tasks.add_task(
+        training_manager.fine_tune, fine_tune_request, job, file_manager
+    )
+    job.status = JobStatus.QUEUED
+    return JobStatusModel(id=job.id, status=job.status.value)
+
+
+# List fine-tune jobs (paginated)
+@router.get("", response_model=PaginatedResponse[JobStatusModel])
+def get_jobs(request: Request):
+    job_manager = request.app.state.job_manager
+    return PaginatedResponse(
+        data=[
+            JobStatusModel(id=job.id, status=job.status.value)
+            for job in job_manager.get_jobs()
+        ],
+        has_more=False,
+    )
+
+
+# List fine-tuning events
+@router.get("/{job_id}/events", response_model=PaginatedResponse[JobEventModel])
+def get_job_events(request: Request, job_id: str):
+    job_manager = request.app.state.job_manager
+    job = job_manager.get_job(job_id)
+    return PaginatedResponse(
+        data=[
+            JobEventModel(
+                id=event.id,
+                level=event.level,
+                message=event.message,
+                data=event.data,
+                created_at=int(event.created_at.timestamp()),
+                type="message",
+            )
+            for event in job.get_events()
+        ],
+        has_more=False,
+    )
+
+
+# List fine-tuning checkpoints
+@router.get(
+    "/{job_id}/checkpoints", response_model=PaginatedResponse[JobCheckpointModel]
+)
+def get_job_checkpoints(request: Request, job_id: str):
+    job_manager = request.app.state.job_manager
+    job = job_manager.get_job(job_id)
+    return PaginatedResponse(
+        data=[
+            JobCheckpointModel(
+                id=checkpoint.id,
+                fine_tuned_model_checkpoint=checkpoint.fine_tuned_model_checkpoint,
+                fine_tuning_job_id=checkpoint.fine_tuning_job_id,
+                metrics=checkpoint.metrics,
+                step_number=checkpoint.step_number,
+            )
+            for checkpoint in job.get_checkpoints()
+        ],
+        has_more=False,
+    )
+
+
+# Retrieve a fine-tune job by id
+@router.get("/{job_id}", response_model=JobStatusModel)
 def get_job_status(
+    request: Request,
     job_id: str,
-    job_manager: JobManager = Depends(get_job_manager)
 ):
+    job_manager = request.app.state.job_manager
     job = job_manager.get_job(job_id)
-    return JobStatusResponse(id=job_id, status=job.status.value, fine_tuned_model=job.fine_tuned_model)
+    return JobStatusModel(
+        id=job_id, status=job.status.value, fine_tuned_model=job.fine_tuned_model
+    )
+
+
+# Cancel a fine-tune job
+@router.post("/{job_id}/cancel", response_model=JobStatusModel)
+def cancel_job(request: Request, job_id: str):
+    job_manager = request.app.state.job_manager
+    job = job_manager.get_job(job_id)
+
+    # Only allow cancellation of jobs that aren't finished
+    if job.status in [JobStatus.SUCCEEDED, JobStatus.FAILED, JobStatus.CANCELLED]:
+        raise HTTPException(
+            status_code=400, detail=f"Cannot cancel job with status {job.status.value}"
+        )
+
+    job.status = JobStatus.PENDING_CANCEL
+    return JobStatusModel(id=job.id, status=job.status.value)

arbor/server/core/config.py

@@ -1,10 +1,47 @@
-from pydantic_settings import BaseSettings
+from pathlib import Path
+from typing import Optional
 
-class Settings(BaseSettings):
-    UPLOADS_DIR: str = "uploads"
-    MODEL_CACHE_DIR: str = "model_cache"
+import yaml
+from pydantic import BaseModel, ConfigDict
 
-    class Config:
-        env_file = ".env"
 
-settings = Settings()
+class InferenceConfig(BaseModel):
+    gpu_ids: str = "0"
+
+
+class TrainingConfig(BaseModel):
+    gpu_ids: str = "0"
+    accelerate_config: Optional[str] = None
+
+
+class ArborConfig(BaseModel):
+    inference: InferenceConfig
+    training: TrainingConfig
+
+
+class Settings(BaseModel):
+
+    STORAGE_PATH: str = "./storage"
+    INACTIVITY_TIMEOUT: int = 30  # 5 seconds
+    arbor_config: ArborConfig
+
+    @classmethod
+    def load_from_yaml(cls, yaml_path: str) -> "Settings":
+        if not yaml_path:
+            raise ValueError("Config file path is required")
+        if not Path(yaml_path).exists():
+            raise ValueError(f"Config file {yaml_path} does not exist")
+
+        try:
+            with open(yaml_path, "r") as f:
+                config = yaml.safe_load(f)
+
+            settings = cls(
+                arbor_config=ArborConfig(
+                    inference=InferenceConfig(**config["inference"]),
+                    training=TrainingConfig(**config["training"]),
+                )
+            )
+            return settings
+        except Exception as e:
+            raise ValueError(f"Error loading config file {yaml_path}: {e}")