diff --git a/examples/finetuning/finetuning_client.ipynb b/examples/finetuning/finetuning_client.ipynb new file mode 100644 index 00000000..5a400ebd --- /dev/null +++ b/examples/finetuning/finetuning_client.ipynb @@ -0,0 +1,140 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Setting up the finetuning job" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "plaintext" + } + }, + "outputs": [], + "source": [ + "from bespokelabs import curator" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "plaintext" + } + }, + "outputs": [], + "source": [ + "finetuning_client = curator.Finetune(\n", + " backend = \"bespoke\",\n", + " backend_params = {\n", + " \"base_url\": 'https://api-dev.bespokelabs.ai'\n", + " }\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "plaintext" + } + }, + "outputs": [], + "source": [ + "dataset_id = \"41da7b9b3e384ae486a9945376d2fd9c\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "plaintext" + } + }, + "outputs": [], + "source": [ + "finetuning_client.create_job(\n", + " model_name=\"Qwen/Qwen2.5-7B-Instruct\",\n", + " job_name=\"bespoke-devrev-demo-live\",\n", + " dataset_id=dataset_id,\n", + " seed=42,\n", + " suffix=\"ft\",\n", + " method = {\n", + " \"type\": \"supervised\",\n", + " \"hyperparameters\": {\n", + " ## Type\n", + " \"finetuning_type\": \"lora\",\n", + " \"lora_rank\": 16,\n", + "\n", + " ## training dynamics\n", + " \"learning_rate\": 0.0001,\n", + " \"num_train_epochs\": 1,\n", + " \"per_device_train_batch_size\": 4,\n", + " \"gradient_accumulation_steps\": 1,\n", + "\n", + " ## infra\n", + " \"preprocessing_num_workers\": 32,\n", + " \"dataloader_num_workers\": 16,\n", + " \"logging_steps\": 1,\n", + " }\n", + " },\n", + " num_gpus=8\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Deploying the finetuned model" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "models_client = curator.Models(\n", + " backend = \"bespoke\",\n", + " backend_params = {\n", + " \"base_url\": 'https://api-dev.bespokelabs.ai'\n", + " }\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "models_client.list_models(job_id='d3201c37255247c2825ff4ce52c110d6')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "plaintext" + } + }, + "outputs": [], + "source": [ + "models_client.deploy_model(model_id = '10426763-76be-483e-987b-07d2ee5adbd7')" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/src/bespokelabs/curator/__init__.py b/src/bespokelabs/curator/__init__.py index 77f0760f..9054ad7f 100644 --- a/src/bespokelabs/curator/__init__.py +++ b/src/bespokelabs/curator/__init__.py @@ -1,10 +1,12 @@ """BespokeLabs Curator.""" from .code_executor.code_executor import CodeExecutor +from .finetune.finetune import Finetune from .llm.llm import LLM +from .models import Models from .types import prompt as types from .utils import load_dataset, push_to_viewer -__all__ = ["LLM", "CodeExecutor", "types", "push_to_viewer", "load_dataset"] +__all__ = ["LLM", "CodeExecutor", "types", "Finetune", "Models", "push_to_viewer", "load_dataset"] from .log import _CONSOLE # noqa: F401 diff --git a/src/bespokelabs/curator/constants.py b/src/bespokelabs/curator/constants.py index 565a95f6..bfbc4650 100644 --- a/src/bespokelabs/curator/constants.py +++ b/src/bespokelabs/curator/constants.py @@ -3,8 +3,8 @@ BATCH_REQUEST_ID_TAG = "custom_id" _CURATOR_DEFAULT_CACHE_DIR = "~/.cache/curator" _DEFAULT_CACHE_DIR = "~/.cache" -BASE_CLIENT_URL = "https://api.bespokelabs.ai/v0/viewer" -PUBLIC_CURATOR_VIEWER_HOME_URL = "https://curator.bespokelabs.ai" +BASE_CLIENT_URL = "https://api-dev.bespokelabs.ai/v0/viewer" +PUBLIC_CURATOR_VIEWER_HOME_URL = "https://curator-dev.bespokelabs.ai" PUBLIC_CURATOR_VIEWER_DATASET_URL = PUBLIC_CURATOR_VIEWER_HOME_URL + "/datasets" _INTERNAL_PROMPT_KEY = "__internal_prompt" _CACHE_MSG = ( diff --git a/src/bespokelabs/curator/datasets/__init__.py b/src/bespokelabs/curator/datasets/__init__.py new file mode 100644 index 00000000..72d10f5a --- /dev/null +++ b/src/bespokelabs/curator/datasets/__init__.py @@ -0,0 +1,5 @@ +"""Dataset loading and processing utilities.""" + +from .base import upload + +__all__ = ["upload"] diff --git a/src/bespokelabs/curator/datasets/base.py b/src/bespokelabs/curator/datasets/base.py new file mode 100644 index 00000000..a3675123 --- /dev/null +++ b/src/bespokelabs/curator/datasets/base.py @@ -0,0 +1,54 @@ +# just upload files, a thin wrapper around push_to_viewer + +import logging +from pathlib import Path + +from datasets import load_dataset + +from bespokelabs.curator.utils import push_to_viewer + +logger = logging.getLogger(__name__) + + +def upload(path: str, split: str = "train"): + """Uploads a dataset file to the Hugging Face Hub. + + Args: + path: The local path to the dataset file. + split: The name of the split to upload (e.g., "train", "test"). + + Raises: + FileNotFoundError: If the specified file does not exist. + ValueError: If the file type is not supported. + """ + # load file into a huggingface dataset + + # it could be a huggingface dataset or a local file + # first check if it is a huggingface dataset + try: + dataset = load_dataset(path, split=split) + + except Exception: + path = Path(path) + if not path.exists(): + raise FileNotFoundError(f"File {path} does not exist") from None + + if path.suffix not in [".jsonl", ".json", ".csv", ".parquet"]: + raise ValueError("Only jsonl, json, csv, and parquet files are supported currently") from None + + try: + if path.suffix == ".jsonl" or path.suffix == ".json": + format = "json" + elif path.suffix == ".csv": + format = "csv" + elif path.suffix == ".parquet": + format = "parquet" + + dataset = load_dataset(format, data_files=str(path), split=split) + + except Exception as e: + logger.error(f"Error loading dataset: {e}") + raise e + + link = push_to_viewer(dataset) + return link diff --git a/src/bespokelabs/curator/finetune/__init__.py b/src/bespokelabs/curator/finetune/__init__.py new file mode 100644 index 00000000..63164cb7 --- /dev/null +++ b/src/bespokelabs/curator/finetune/__init__.py @@ -0,0 +1,5 @@ +"""Finetuning backend abstraction and implementations.""" + +from .finetune import Finetune + +__all__ = ["Finetune"] diff --git a/src/bespokelabs/curator/finetune/base_backend.py b/src/bespokelabs/curator/finetune/base_backend.py new file mode 100644 index 00000000..f17aa621 --- /dev/null +++ b/src/bespokelabs/curator/finetune/base_backend.py @@ -0,0 +1,58 @@ +from abc import ABC, abstractmethod + + +class BaseFinetuneBackend(ABC): + """Abstract base class for finetuning backends.""" + + def __init__(self, backend_params: dict): + """Initializes the finetuning backend. + + Args: + backend_params: A dictionary containing backend-specific parameters. + """ + self.backend_params = backend_params + + @abstractmethod + def create_job(self, *args, **kwargs): + """Creates a new finetuning job.""" + pass + + @abstractmethod + def list_jobs(self, *args, **kwargs): + """Lists all finetuning jobs.""" + pass + + @abstractmethod + def list_job_events(self, job_id: str): + """Lists events for a specific finetuning job. + + Args: + job_id: The ID of the finetuning job. + """ + pass + + # @abstractmethod + # def list_job_checkpoints(self, job_id: str): + # pass + + # @abstractmethod + # def list_job_metrics(self, job_id: str): + # pass + + @abstractmethod + def get_job_details(self, job_id: str): + """Retrieves details for a specific finetuning job. + + Args: + job_id: The ID of the finetuning job. + """ + pass + + @abstractmethod + def cancel_job(self, job_id: str): + """Cancels a specific finetuning job. + + Args: + job_id: The ID of the finetuning job. + """ + pass diff --git a/src/bespokelabs/curator/finetune/bespoke_backend.py b/src/bespokelabs/curator/finetune/bespoke_backend.py new file mode 100644 index 00000000..d99d875b --- /dev/null +++ b/src/bespokelabs/curator/finetune/bespoke_backend.py @@ -0,0 +1,147 @@ +import logging +import os +from typing import Union + +import requests + +from bespokelabs.curator.finetune.base_backend import BaseFinetuneBackend +from bespokelabs.curator.finetune.types import FinetuneBackendParams + +logger = logging.getLogger(__name__) + + +class BespokeFinetuneBackend(BaseFinetuneBackend): + """A client for interacting with the Bespoke Labs Finetuning API. + + This class provides methods to create, list, and manage finetuning jobs + on the Bespoke Labs platform. + """ + + def __init__(self, backend_params: Union[FinetuneBackendParams, dict]): + """Initializes the BespokeFinetuneBackend. + + Args: + backend_params: A dictionary or FinetuneBackendParams object containing + configuration for the backend, such as base_url and env. + """ + super().__init__(backend_params) + # Determine the environment suffix for the API URL. + self.env = "-dev" if backend_params.get("env", "prod") == "dev" else "" + # Set the base URL for the finetuning API. + if not backend_params.get("base_url"): + self.base_url = f"https://api{self.env}.bespokelabs.com/v0/finetune" + else: + self.base_url = backend_params["base_url"] + "/v0/finetune" + # Retrieve the API key from environment variables. + self.api_key = os.environ.get("BESPOKE_API_KEY") + + def create_job(self, *args, **kwargs): + """Creates a new finetuning job. + + Args: + *args: Variable length argument list. + **kwargs: Arbitrary keyword arguments. Expected keys include: + 'model_name': The name of the model to finetune. + 'dataset_id': The ID of the dataset for training. + 'seed': The random seed for reproducibility. + 'suffix': A suffix to append to the finetuned model name. + 'method': The finetuning method and its hyperparameters. + 'job_name': A name for the finetuning job. + 'num_gpus': The number of GPUs to use for finetuning. + + Returns: + dict: The JSON response from the API, typically containing job details. + """ + url = f"{self.base_url}/jobs/create" + headers = {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"} + + # Warn if Hugging Face token is not set, as it might be needed for private resources. + if os.environ.get("HF_TOKEN") is None: + logger.warning("HF_TOKEN is not set. Please set it to use private gated models or datasets.") + + # Prepare the data payload for the API request. + data = { + "model_name": kwargs["model_name"], + "dataset_id": kwargs["dataset_id"], + "seed": kwargs["seed"], + "suffix": kwargs["suffix"], + "method": kwargs["method"], + "job_name": kwargs["job_name"], + "num_gpus": kwargs["num_gpus"], + "secrets": { + "HF_TOKEN": os.environ.get("HF_TOKEN", None), # Include HF_TOKEN if available. + }, + } + + response = requests.post(url, headers=headers, json=data) + return response.json() + + def list_jobs(self, *args, **kwargs): + """Lists all finetuning jobs. + + Args: + *args: Variable length argument list (currently unused). + **kwargs: Arbitrary keyword arguments (currently unused). + + Returns: + dict: The JSON response from the API, typically a list of jobs. + """ + url = f"{self.base_url}/jobs" + headers = {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"} + + response = requests.get(url, headers=headers) + return response.json() + + def list_job_events(self, *args, **kwargs): + """Lists events for a specific finetuning job. + + Args: + *args: Variable length argument list (currently unused). + **kwargs: Arbitrary keyword arguments. Expected key: + 'job_id': The ID of the job to retrieve events for. + + Returns: + dict: The JSON response from the API, typically a list of job events. + """ + url = f"{self.base_url}/jobs/{kwargs['job_id']}/events" + headers = {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"} + + response = requests.get(url, headers=headers) + return response.json() + + def get_job_details(self, *args, **kwargs): + """Retrieves details for a specific finetuning job. + + Args: + *args: Variable length argument list (currently unused). + **kwargs: Arbitrary keyword arguments. Expected keys: + 'job_id': The ID of the job to retrieve details for. + 'type' (optional): The type of logs to retrieve (e.g., 'EVENTS'). + Defaults to 'EVENTS'. + + Returns: + dict: The JSON response from the API, containing job details. + """ + # Construct URL with optional log_type parameter. + url = f"{self.base_url}/jobs/{kwargs['job_id']}?log_type={kwargs.get('type', 'EVENTS')}" + headers = {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"} + + response = requests.get(url, headers=headers) + return response.json() + + def cancel_job(self, *args, **kwargs): + """Cancels a specific finetuning job. + + Args: + *args: Variable length argument list (currently unused). + **kwargs: Arbitrary keyword arguments. Expected key: + 'job_id': The ID of the job to cancel. + + Returns: + dict: The JSON response from the API, typically confirming cancellation. + """ + url = f"{self.base_url}/jobs/{kwargs['job_id']}/cancel" + headers = {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"} + + response = requests.post(url, headers=headers) + return response.json() diff --git a/src/bespokelabs/curator/finetune/finetune.py b/src/bespokelabs/curator/finetune/finetune.py new file mode 100644 index 00000000..8c354277 --- /dev/null +++ b/src/bespokelabs/curator/finetune/finetune.py @@ -0,0 +1,47 @@ +from typing import Union + +from bespokelabs.curator.finetune.bespoke_backend import BespokeFinetuneBackend +from bespokelabs.curator.finetune.openai_backend import OpenAIFinetuneBackend +from bespokelabs.curator.finetune.types import FinetuneBackendParams + + +class Finetune: + """A class to manage finetuning jobs using different backends.""" + + def __init__(self, backend: str, backend_params: Union[FinetuneBackendParams, dict]): + """Initializes the Finetune class with a specified backend. + + Args: + backend: The name of the backend to use (e.g., "bespoke", "openai"). + backend_params: Parameters for the backend. + + Raises: + ValueError: If an invalid backend is specified. + """ + self.backend_params = backend_params + if backend == "bespoke": + self._backend = BespokeFinetuneBackend(backend_params) + elif backend == "openai": + self._backend = OpenAIFinetuneBackend(backend_params) + else: + raise ValueError(f"Invalid backend: {backend}") + + def create_job(self, *args, **kwargs): + """Creates a finetuning job using the configured backend.""" + return self._backend.create_job(*args, **kwargs) + + def list_jobs(self, *args, **kwargs): + """Lists finetuning jobs using the configured backend.""" + return self._backend.list_jobs(*args, **kwargs) + + def list_job_events(self, *args, **kwargs): + """Lists job events using the configured backend.""" + return self._backend.list_job_events(*args, **kwargs) + + def get_job_details(self, *args, **kwargs): + """Gets job details using the configured backend.""" + return self._backend.get_job_details(*args, **kwargs) + + def cancel_job(self, *args, **kwargs): + """Cancels a job using the configured backend.""" + return self._backend.cancel_job(*args, **kwargs) diff --git a/src/bespokelabs/curator/finetune/openai_backend.py b/src/bespokelabs/curator/finetune/openai_backend.py new file mode 100644 index 00000000..b888ab2b --- /dev/null +++ b/src/bespokelabs/curator/finetune/openai_backend.py @@ -0,0 +1,139 @@ +import os + +from openai import OpenAI + +from bespokelabs.curator.finetune.base_backend import BaseFinetuneBackend + + +class OpenAIFinetuneBackend(BaseFinetuneBackend): + """A client for interacting with the OpenAI Finetuning API. + + This class provides methods to create, list, and manage finetuning jobs + on the OpenAI platform. It serves as a backend for the Finetune class. + """ + + def __init__(self, backend_params: dict): + """Initializes the OpenAIFinetuneBackend. + + Args: + backend_params: A dictionary containing configuration for the backend. + Currently, this is not used by the OpenAI backend but is + part of the BaseFinetuneBackend interface. + The OpenAI API key is retrieved from the environment variable + "OPENAI_API_KEY". + """ + super().__init__(backend_params) + # Initialize the OpenAI client using the API key from environment variables. + self.client = OpenAI(api_key=os.environ["OPENAI_API_KEY"]) + + def create_job(self, **kwargs): + """Creates a new finetuning job on the OpenAI platform. + + Args: + **kwargs: Arbitrary keyword arguments passed directly to the + `openai.fine_tuning.jobs.create` method. Expected keys include: + 'dataset_id' (str): The ID of an uploaded file that contains training data. + 'model_name' (str): The name of the model to fine-tune. + 'method' (Optional[dict]): The finetuning method and its hyperparameters. + (Note: OpenAI API might use 'hyperparameters' directly + or a nested structure depending on the API version and model type). + 'hyperparameters' (Optional[dict]): The hyperparameters used for the fine-tuning job. + (This might be deprecated in favor of 'method' in some contexts). + 'seed' (Optional[int]): The seed to use for reproducibility. + 'suffix' (Optional[str]): A string of up to 40 characters that will be added to your + fine-tuned model name. + 'validation_file' (Optional[str]): The ID of an uploaded file that contains validation data. + Refer to the OpenAI API documentation for a complete list of parameters. + + Returns: + openai.types.fine_tuning.FineTuningJob: An object representing the created fine-tuning job. + """ + job = self.client.fine_tuning.jobs.create( + training_file=kwargs["dataset_id"], + model=kwargs["model_name"], + hyperparameters=kwargs.get("hyperparameters"), + seed=kwargs.get("seed"), + suffix=kwargs.get("suffix"), + validation_file=kwargs.get("validation_file"), + ) + return job + + def list_jobs(self, **kwargs): + """Lists all finetuning jobs for the organization on the OpenAI platform. + + Args: + **kwargs: Arbitrary keyword arguments passed directly to the + `openai.fine_tuning.jobs.list` method. Common keys include: + 'after' (Optional[str]): Identifier for the last job from the previous pagination request. + 'limit' (Optional[int]): Number of fine-tuning jobs to retrieve. + Refer to the OpenAI API documentation for a complete list of parameters. + + Returns: + openai.pagination.SyncCursorPage[openai.types.fine_tuning.FineTuningJob]: + A paginated list of fine-tuning jobs. + """ + # List fine-tuning jobs using the OpenAI client. + # All provided keyword arguments are passed to the OpenAI API. + jobs = self.client.fine_tuning.jobs.list(**kwargs) + return jobs + + def list_job_events(self, **kwargs): + """Lists events for a specific finetuning job on the OpenAI platform. + + Args: + **kwargs: Arbitrary keyword arguments. Expected key: + 'job_id' (str): The ID of the fine-tuning job to retrieve events for. + This is passed as `fine_tuning_job_id` to the OpenAI API. + Other kwargs are passed directly to `openai.fine_tuning.jobs.list_events`. + Common keys include: + 'after' (Optional[str]): Identifier for the last event from the previous pagination request. + 'limit' (Optional[int]): Number of events to retrieve. + Refer to the OpenAI API documentation for a complete list of parameters. + + Returns: + openai.pagination.SyncCursorPage[openai.types.fine_tuning.FineTuningJobEvent]: + A paginated list of fine-tuning job events. + """ + # Retrieve the job_id from kwargs and remove it to prevent it from being passed again. + fine_tuning_job_id = kwargs.pop("job_id") + # List events for a specific fine-tuning job using the OpenAI client. + events = self.client.fine_tuning.jobs.list_events(fine_tuning_job_id=fine_tuning_job_id, **kwargs) + return events + + def get_job_details(self, **kwargs): + """Retrieves details for a specific finetuning job on the OpenAI platform. + + Args: + **kwargs: Arbitrary keyword arguments. Expected key: + 'job_id' (str): The ID of the fine-tuning job to retrieve. + This is passed as `fine_tuning_job_id` to the OpenAI API. + Other kwargs are passed directly to `openai.fine_tuning.jobs.retrieve`. + Refer to the OpenAI API documentation for a complete list of parameters. + + Returns: + openai.types.fine_tuning.FineTuningJob: An object representing the fine-tuning job details. + """ + # Retrieve the job_id from kwargs and remove it to prevent it from being passed again. + fine_tuning_job_id = kwargs.pop("job_id") + # Retrieve details for a specific fine-tuning job using the OpenAI client. + job = self.client.fine_tuning.jobs.retrieve(fine_tuning_job_id=fine_tuning_job_id, **kwargs) + return job + + def cancel_job(self, **kwargs): + """Cancels a specific finetuning job on the OpenAI platform. + + Args: + **kwargs: Arbitrary keyword arguments. Expected key: + 'job_id' (str): The ID of the fine-tuning job to cancel. + This is passed as `fine_tuning_job_id` to the OpenAI API. + Other kwargs are passed directly to `openai.fine_tuning.jobs.cancel`. + Refer to the OpenAI API documentation for a complete list of parameters. + + Returns: + openai.types.fine_tuning.FineTuningJob: An object representing the cancelled fine-tuning job. + """ + # Retrieve the job_id from kwargs and remove it to prevent it from being passed again. + fine_tuning_job_id = kwargs.pop("job_id") + # Cancel a specific fine-tuning job using the OpenAI client. + job = self.client.fine_tuning.jobs.cancel(fine_tuning_job_id=fine_tuning_job_id, **kwargs) + return job diff --git a/src/bespokelabs/curator/finetune/types.py b/src/bespokelabs/curator/finetune/types.py new file mode 100644 index 00000000..26cc6da7 --- /dev/null +++ b/src/bespokelabs/curator/finetune/types.py @@ -0,0 +1,50 @@ +from typing import Optional + +from pydantic import BaseModel + + +class FinetuneMethod(BaseModel): + """Pydantic model for specifying the finetuning method and its hyperparameters.""" + + type: str + hyperparameters: dict + + +class FinetuneRequest(BaseModel): + """Pydantic model for a finetuning job request.""" + + job_name: str + dataset_id: str + model_name: str + seed: int + suffix: str + num_gpus: int + method: FinetuneMethod + + +class FinetuneResponse(BaseModel): + """Pydantic model for a finetuning job response.""" + + job_id: str + job_name: str + dataset_id: str + model_name: str + seed: int + suffix: str + method: FinetuneMethod + status: str + + +class FinetuneBackendParams(BaseModel): + """Pydantic model for backend parameters.""" + + base_url: Optional[str] = None + api_key: Optional[str] = None + model_name: Optional[str] = None + dataset_id: Optional[str] = None + method: Optional[FinetuneMethod] = None + seed: Optional[int] = None + suffix: Optional[str] = None + num_gpus: Optional[int] = None + hyperparameters: Optional[dict] = None + validation_file: Optional[str] = None diff --git a/src/bespokelabs/curator/models/__init__.py b/src/bespokelabs/curator/models/__init__.py new file mode 100644 index 00000000..7ef3a2dd --- /dev/null +++ b/src/bespokelabs/curator/models/__init__.py @@ -0,0 +1,5 @@ +"""Models module for Curator.""" + +from .models import Models + +__all__ = ["Models"] diff --git a/src/bespokelabs/curator/models/base_backend.py b/src/bespokelabs/curator/models/base_backend.py new file mode 100644 index 00000000..46faa648 --- /dev/null +++ b/src/bespokelabs/curator/models/base_backend.py @@ -0,0 +1,21 @@ +import typing as t +from abc import ABC, abstractmethod + + +class BaseModelsBackend(ABC): + """Base interface for model operations.""" + + @abstractmethod + def list_models(self, job_id: str = None) -> t.List[dict]: + """List available models or models related to a specific job.""" + pass + + @abstractmethod + def deploy_model(self, model_id: str) -> bool: + """Deploy a model to make it available for inference.""" + pass + + @abstractmethod + def undeploy_model(self, model_id: str) -> bool: + """Undeploy a model to free up resources.""" + pass diff --git a/src/bespokelabs/curator/models/bespoke_backend.py b/src/bespokelabs/curator/models/bespoke_backend.py new file mode 100644 index 00000000..7ae53c8b --- /dev/null +++ b/src/bespokelabs/curator/models/bespoke_backend.py @@ -0,0 +1,88 @@ +import os +import typing as t + +import requests + +from ..log import logger +from .base_backend import BaseModelsBackend + + +class BespokeModelsBackend(BaseModelsBackend): + """Client for interacting with Bespoke models API.""" + + def __init__(self, base_url: str = None, api_key: str = None): + """Initialize the Bespoke models client. + + Args: + base_url: Base URL for the API + api_key: API key for authentication (optional) + """ + self.base_url = base_url + self.api_key = api_key or os.environ.get("BESPOKE_API_KEY") + self.headers = {"Authorization": f"Bearer {self.api_key}"} if self.api_key else {} + + def list_models(self, job_id: str = None) -> t.List[dict]: + """List available models or models related to a specific job. + + Args: + job_id: Optional job ID to filter models by + + Returns: + A list of model information dictionaries + """ + url = f"{self.base_url}/v0/models/{job_id}" + + try: + response = requests.get(url, headers=self.headers) + if response.status_code == 200: + return response.json() + else: + logger.error(f"Failed to list models: {response.status_code}, {response.text}") + return [] + except Exception as e: + logger.error(f"Error listing models: {str(e)}") + return [] + + def deploy_model(self, model_id: str) -> bool: + """Deploy a model to make it available for inference. + + Args: + model_id: ID of the model to deploy + + Returns: + True if deployment was successful, False otherwise + """ + url = f"{self.base_url}/v0/models/{model_id}/deploy" + + try: + response = requests.post(url, headers=self.headers) + if response.status_code in [200, 201, 202]: + return response.json() + else: + logger.error(f"Failed to deploy model: {response.status_code}, {response.text}") + return False + except Exception as e: + logger.error(f"Error deploying model: {str(e)}") + return False + + def undeploy_model(self, model_id: str) -> bool: + """Undeploy a model to free up resources. + + Args: + model_id: ID of the model to undeploy + + Returns: + True if undeployment was successful, False otherwise + """ + url = f"{self.base_url}/v0/models/{model_id}/undeploy" + + try: + response = requests.post(url, headers=self.headers) + if response.status_code in [200, 202, 204]: + return response.json() + else: + logger.error(f"Failed to undeploy model: {response.status_code}, {response.text}") + return False + except Exception as e: + logger.error(f"Error undeploying model: {str(e)}") + return False diff --git a/src/bespokelabs/curator/models/models.py b/src/bespokelabs/curator/models/models.py new file mode 100644 index 00000000..af5f6037 --- /dev/null +++ b/src/bespokelabs/curator/models/models.py @@ -0,0 +1,34 @@ +import typing as t + +from .bespoke_backend import BespokeModelsBackend + + +class Models: + """Factory for creating model clients based on the specified backend.""" + + def __init__(self, backend: str = "bespoke", backend_params: dict = None): + """Initialize the Models factory. + + Args: + backend: The backend provider to use ('bespoke' supported) + backend_params: Configuration parameters for the backend + """ + self.backend = backend + self.backend_params = backend_params or {} + + if backend == "bespoke": + self.client = BespokeModelsBackend(base_url=self.backend_params.get("base_url"), api_key=self.backend_params.get("api_key")) + else: + raise ValueError(f"Unsupported backend: {backend}") + + def list_models(self, job_id: str = None) -> t.List[dict]: + """Delegate to the client's list_models method.""" + return self.client.list_models(job_id=job_id) + + def deploy_model(self, model_id: str) -> bool: + """Delegate to the client's deploy_model method.""" + return self.client.deploy_model(model_id=model_id) + + def undeploy_model(self, model_id: str) -> bool: + """Delegate to the client's undeploy_model method.""" + return self.client.undeploy_model(model_id=model_id) diff --git a/src/bespokelabs/curator/utils.py b/src/bespokelabs/curator/utils.py index 869dabc9..9fea6301 100644 --- a/src/bespokelabs/curator/utils.py +++ b/src/bespokelabs/curator/utils.py @@ -98,6 +98,7 @@ async def send_row(idx, row): await client.session_completed() run_in_event_loop(send_responses()) + return view_url