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

arbor/client/api.py

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

arbor/server/__init__.py

@@ -0,0 +1 @@
+

arbor/server/api/__init__.py

@@ -0,0 +1 @@
+

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/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/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

@@ -0,0 +1,117 @@
+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()
+
+
+# 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 = request.app.state.job_manager
+    job = job_manager.get_job(job_id)
+    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/__init__.py

@@ -0,0 +1 @@
+

arbor/server/core/config.py

@@ -0,0 +1,47 @@
+from pathlib import Path
+from typing import Optional
+
+import yaml
+from pydantic import BaseModel, ConfigDict
+
+
+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}")

arbor/server/main.py

@@ -0,0 +1,11 @@
+from fastapi import FastAPI
+
+from arbor.server.api.routes import files, grpo, inference, jobs
+
+app = FastAPI(title="Arbor API")
+
+# Include routers
+app.include_router(files.router, prefix="/v1/files")
+app.include_router(jobs.router, prefix="/v1/fine_tuning/jobs")
+app.include_router(grpo.router, prefix="/v1/fine_tuning/grpo")
+app.include_router(inference.router, prefix="/v1/chat")