openweights 0.1.0__py3-none-any.whl

openweights/__init__.py

@@ -0,0 +1,3 @@
+from .client import OpenWeights
+
+__all__ = ['OpenWeights']

openweights/client/__init__.py

@@ -0,0 +1,176 @@
+import asyncio
+import atexit
+import json
+from typing import Optional, BinaryIO, Dict, Any, List, Union
+import os
+import sys
+from postgrest.exceptions import APIError
+import hashlib
+from datetime import datetime
+from openai import OpenAI, AsyncOpenAI
+import backoff
+import time
+from supabase import create_client, Client
+from supabase.lib.client_options import ClientOptions
+
+from openweights.validate import validate_messages, validate_preference_dataset, TrainingConfig, InferenceConfig, ApiConfig
+from openweights.client.files import Files
+from openweights.client.jobs import FineTuningJobs, InferenceJobs, Deployments, Jobs
+from openweights.client.run import Run, Runs
+from openweights.client.events import Events
+from openweights.client.temporary_api import TemporaryApi, group_models_or_adapters_by_model, get_lora_rank
+from openweights.client.chat import ChatCompletions, AsyncChatCompletions
+from openweights.client.custom_job import CustomJob
+
+
+def create_authenticated_client(supabase_url: str, supabase_anon_key: str, auth_token: Optional[str] = None):
+    """Create a Supabase client with authentication.
+    
+    Args:
+        supabase_url: Supabase project URL
+        supabase_anon_key: Supabase anon key
+        auth_token: Session token from Supabase auth (optional)
+        api_key: OpenWeights API key starting with 'ow_' (optional)
+    """
+    headers = {}
+    if auth_token:
+        headers["Authorization"] = f"Bearer {auth_token}"
+    else:
+        raise ValueError("No auth_token provided")
+    
+    options = ClientOptions(
+        schema="public",
+        headers=headers,
+        auto_refresh_token=False,
+        persist_session=False
+    )
+    
+    return create_client(supabase_url, supabase_anon_key, options)
+
+
+class OpenWeights:
+    def __init__(self, 
+                 supabase_url: Optional[str] = None, 
+                 supabase_key: Optional[str] = None, 
+                 auth_token: Optional[str] = None,
+                 organization_id: Optional[str] = None,
+                 use_async: bool = False,
+                 deploy_kwargs: Dict[str, Any] = {'max_model_len': 2048}):
+        """Initialize OpenWeights client
+        
+        Args:
+            supabase_url: Supabase project URL (or SUPABASE_URL env var)
+            supabase_key: Supabase anon key (or SUPABASE_ANON_KEY env var)
+            auth_token: Authentication token (or OPENWEIGHTS_API_KEY env var)
+                       Can be either a session token or a service account JWT token
+        """
+        self.supabase_url = supabase_url or os.environ.get('SUPABASE_URL', 'https://taofkfabrhpgtohaikst.supabase.co')
+        self.supabase_key = supabase_key or os.environ.get('SUPABASE_ANON_KEY', 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6InRhb2ZrZmFicmhwZ3RvaGFpa3N0Iiwicm9sZSI6ImFub24iLCJpYXQiOjE3MzE5MjkyMjcsImV4cCI6MjA0NzUwNTIyN30.KRufleTgprt16mfm0_91YjKIFZAne1-IW8buMxWVMeE')
+        self.auth_token = auth_token or os.getenv('OPENWEIGHTS_API_KEY')
+        self.deploy_kwargs = deploy_kwargs
+        
+        if not self.supabase_url or not self.supabase_key:
+            raise ValueError("Supabase URL and key must be provided either as arguments or environment variables")
+        
+        if not self.auth_token:
+            raise ValueError("Authentication token must be provided either as argument or OPENWEIGHTS_API_KEY environment variable")
+        
+        self._supabase = create_authenticated_client(
+            self.supabase_url, 
+            self.supabase_key, 
+            self.auth_token
+        )
+        
+        # Get organization ID from token
+        self.organization_id = organization_id or self.get_organization_id()
+        self.org_name = self.get_organization_name()
+        print("Connected to org: ", self.org_name)
+        self.set_hf_org_env()
+        
+        # Initialize components with organization ID
+        self.files = Files(self._supabase, self.organization_id)
+        self.fine_tuning = FineTuningJobs(self._supabase, self.organization_id)
+        self.inference = InferenceJobs(self._supabase, self.organization_id)
+        self.jobs = Jobs(self._supabase, self.organization_id)
+        self.deployments = Deployments(self._supabase, self.organization_id)
+        self.runs = Runs(self._supabase)
+        self.events = Events(self._supabase)
+        self.async_chat = AsyncChatCompletions(self, deploy_kwargs=self.deploy_kwargs)
+        self.sync_chat = ChatCompletions(self, deploy_kwargs=self.deploy_kwargs)
+        self.chat = self.async_chat if use_async else self.sync_chat
+
+        self._current_run = None
+    
+    @backoff.on_exception(backoff.constant, Exception, interval=1, max_time=60, max_tries=60, on_backoff=lambda details: print(f"Retrying... {details['exception']}"))
+    def get_organization_id(self) -> str:
+        """Get the organization ID associated with the current token"""
+        result = self._supabase.rpc('get_organization_from_token').execute()
+        if not result.data:
+            raise ValueError("Could not determine organization ID from token")
+        return result.data
+    
+    @backoff.on_exception(backoff.constant, Exception, interval=1, max_time=60, max_tries=60, on_backoff=lambda details: print(f"Retrying... {details['exception']}"))
+    def get_organization_name(self):
+        """Get the organization ID associated with the current token"""
+        result = self._supabase.table('organizations')\
+            .select('*')\
+            .eq('id', self.organization_id)\
+            .single().execute()
+        return result.data['name']
+    
+    @backoff.on_exception(backoff.constant, Exception, interval=1, max_time=60, max_tries=60, on_backoff=lambda details: print(f"Retrying... {details['exception']}"))
+    def set_hf_org_env(self):
+        """Get organization secrets from the database."""
+        if os.environ.get('HF_ORG'):
+            return
+        result = self._supabase.table('organization_secrets')\
+            .select('value')\
+            .eq('organization_id', self.organization_id)\
+            .eq('name', 'HF_ORG')\
+            .single().execute()
+        if not result.data:
+            raise ValueError("Could not determine organization ID from token")
+        os.environ['HF_ORG'] = result.data['value']
+        print(f"Set HF_ORG to {os.environ['HF_ORG']}")
+        
+    @property
+    def run(self):
+        if not self._current_run:
+            self._current_run = Run(self._supabase, organization_id=self.organization_id)
+        return self._current_run
+    
+    def deploy(self, model: str, lora_adapters: List[str] = None, max_lora_rank: str = 'guess', max_model_len: int = 2048, api_key: str = os.environ.get('OW_DEFAULT_API_KEY'), requires_vram_gb: str = 'guess', max_num_seqs: int = 100) -> TemporaryApi:
+        if lora_adapters is None:
+            lora_adapters = []
+        """Deploy a model on OpenWeights"""
+        if api_key is None:
+            api_key = self.auth_token
+        if lora_adapters and max_lora_rank == 'guess':
+            max_lora_rank = max(get_lora_rank(a) for a in lora_adapters)
+        else:
+            max_lora_rank = 16
+        job = self.deployments.create(
+            model=model, max_model_len=max_model_len, api_key=api_key, requires_vram_gb=requires_vram_gb,
+            lora_adapters=lora_adapters, max_lora_rank=max_lora_rank, max_num_seqs=max_num_seqs)
+        return TemporaryApi(self, job['id'])
+    
+    def multi_deploy(self, models: List[str], max_model_len: Union[int,str] = 2048, api_key: str = os.environ.get('OW_DEFAULT_API_KEY'), requires_vram_gb: Union[int,str] = 'guess', max_num_seqs: int = 100, base_model_override: Optional[str] = None) -> Dict[str, TemporaryApi]:
+        """Deploy multiple models - creates on server for each base model, and deploys all lora adapters on of the same base model together"""
+        assert isinstance(models, list), "models must be a list"
+        lora_groups = group_models_or_adapters_by_model(models)
+        apis = {}
+        for model, lora_adapters in lora_groups.items():
+            if base_model_override is not None:
+                model = base_model_override
+            print(f"Deploying {model} with {len(lora_adapters)} lora adapters")
+            api = self.deploy(model, lora_adapters=lora_adapters, max_model_len=max_model_len, api_key=api_key, requires_vram_gb=requires_vram_gb, max_num_seqs=max_num_seqs)
+            for model_id in [model] + lora_adapters:
+                apis[model_id] = api
+        return apis
+    
+    def register(self, name: str):
+        """Decorator to register a custom job class"""
+        def register_custom_job(cls):
+            obj = cls(self)
+            setattr(self, name, obj)
+        return register_custom_job

openweights/client/cache_on_disk.py

@@ -0,0 +1,34 @@
+import hashlib
+import json
+import os
+from functools import wraps
+
+import diskcache as dc
+
+# Ensure the cache directory exists
+cache_dir = os.path.join(os.path.dirname(__file__), ".llm-cache")
+os.makedirs(cache_dir, exist_ok=True)
+
+# Use FanoutCache with multiple shards and a short timeout
+cache = dc.FanoutCache(cache_dir, shards=64, timeout=1)
+
+def cache_on_disk(required_kwargs=[]):
+    def decorator(function):
+        @wraps(function)
+        async def wrapper(*args, **kwargs):
+            # Only cache if all required kwargs are present
+            if not all(k in kwargs for k in required_kwargs):
+                return await function(*args, **kwargs)
+            
+            # Serialize args/kwargs into a JSON string and hash it
+            serialized = json.dumps({"args": args, "kwargs": kwargs}, sort_keys=True)
+            key = hashlib.sha256(serialized.encode()).hexdigest()
+            
+            if key in cache:
+                return cache[key]
+            
+            result = await function(*args, **kwargs)
+            cache[key] = result
+            return result
+        return wrapper
+    return decorator

openweights/client/chat.py

@@ -0,0 +1,135 @@
+from collections import defaultdict
+import openai
+import asyncio
+from openweights.client.cache_on_disk import cache_on_disk
+import backoff
+
+APIS = {}
+DEPLOYMENT_QUEUE = []
+STARTING = []
+
+
+class AsyncChatCompletions:
+    """This class is a wrapper around the OpenAI Chat API that handles deployment of models,
+    request caching (when seeds are provided), and rate limiting.
+    
+    Args:
+        ow: OpenWeights client
+        deploy_kwargs: kwargs to pass to ow.multi_deploy
+        request_timeout: timeout for requests in seconds
+        per_token_timeout: computes a timeout based on max_tokens * per_token_timeout for each request
+
+        When both timeouts are set, the maximum of the two is used.
+    """
+    def __init__(self, ow, deploy_kwargs={}, request_timeout=5, per_token_timeout=1):
+        self.ow = ow
+        self.completions = self
+        self.deploy_kwargs = deploy_kwargs
+        self.request_timeout = request_timeout
+        self.per_token_timeout = per_token_timeout
+    
+    async def create(self, model: str, **kwargs):
+        @cache_on_disk(required_kwargs=['seed'])
+        async def cached_create(model, **kwargs):
+            return await self._create(model, **kwargs)
+        return await cached_create(model, **kwargs)
+    
+    async def _create(self, model, **kwargs):
+        api = await self._get_api(model)
+        async with api.sem:
+            return await self._create_with_backoff(api, model, **kwargs)
+
+    @backoff.on_exception(
+        wait_gen=backoff.expo,
+        exception=(
+            openai.RateLimitError,
+            openai.APIConnectionError,
+            openai.APITimeoutError
+        ),
+        max_value=60,
+        factor=1.5,
+        max_tries=10
+    )
+    async def _create_with_backoff(self, api, model, **kwargs):
+        timeout = kwargs.pop('timeout', None) or max(
+            self.request_timeout,
+            kwargs.get('max_tokens', 1) * self.per_token_timeout,
+        )
+        return await api.async_client.chat.completions.create(model=model, timeout=timeout, **kwargs)
+    
+    async def _get_api(self, model):
+        """If the model is not yet deployed, we add it to a queue of to-be-deployed models and wait for 5 seconds
+        to group them at once, which is more efficient if the multiple lora finetunes of the same model are deployed."""
+        if model in APIS:
+            return APIS[model]
+        if looks_like_openai(model):
+            return OpenAiApi()
+        if model not in DEPLOYMENT_QUEUE and model not in STARTING:
+            print(f"Adding {model} to deployment queue")
+            DEPLOYMENT_QUEUE.append(model)
+            # Create a task to deploy the model in 5 seconds
+            asyncio.create_task(self._wait_and_deploy_queue())
+        # Poll until model is deployed
+        while model not in APIS:
+            await asyncio.sleep(1)
+        return APIS[model]
+    
+    async def _wait_and_deploy_queue(self, seconds_to_wait=5):
+        for _ in range(seconds_to_wait):
+            await asyncio.sleep(1)
+            if len(DEPLOYMENT_QUEUE) == 0:
+                return
+        # Move all models from the queue to a list of models to be deployed, such that DEPLOYMENT_QUEUE is empty after this
+        models_to_deploy = DEPLOYMENT_QUEUE.copy()
+        DEPLOYMENT_QUEUE.clear()
+        STARTING.extend(models_to_deploy)
+        print(f"Deploying {models_to_deploy}")
+        # Deploy the models
+        apis = self.ow.multi_deploy(models_to_deploy, **self.deploy_kwargs)
+        # Wait for apis to be up and move each model to APIS as soon as its API is ready
+        print(f"Waiting for {models_to_deploy} to be up")
+        api_to_models = defaultdict(list)
+        for model, api in apis.items():
+            api_to_models[api].append(model)
+        
+        async def handle_api(api, models):
+            await api.async_up()
+            for model in models:
+                APIS[model] = api
+                STARTING.remove(model)
+        
+        await asyncio.gather(*[handle_api(api, models) for api, models in api_to_models.items()])
+    
+    def kill(self, model_id):
+        api = APIS.pop(model_id, None)
+        if api is not None and api not in APIS.values():
+            api.down()
+        self.ow.jobs.cancel(api.job_id)
+        
+
+
+class ChatCompletions(AsyncChatCompletions):
+    def create(self, **kwargs):
+        assert kwargs.get('stream', False) is False, "ow.chat.completions.create does only support stream=True in async mode"
+        response = asyncio.run(super().create(**kwargs))
+        return response
+
+from openweights.client.temporary_api import TemporaryApi
+import openai
+
+def looks_like_openai(model):
+    return any(model.lower().startswith(i) for i in  ['gpt', 'o1', 'o3'])
+
+class OpenAiApi(TemporaryApi):
+    def __init__(self, concurrents=10, base_url=None, api_key=None, models=[]):
+        self.concurrents = concurrents
+        self.sem = asyncio.Semaphore(concurrents)
+        if base_url is None:
+            self.async_client = openai.AsyncOpenAI()
+            self.sync_client = openai.OpenAI()
+        else:
+            self.async_client = openai.AsyncOpenAI(base_url=base_url, api_key=api_key)
+            self.sync_client = openai.OpenAI(base_url=base_url, api_key=api_key)
+        
+        for model in models:
+            APIS[model] = self

openweights/client/custom_job.py

@@ -0,0 +1,96 @@
+import os
+import json
+from typing import Dict, Optional, Type, Any
+from pydantic import BaseModel, Field
+
+from openweights.client.jobs import BaseJob
+
+
+class CustomJob(BaseJob):
+    """Base class for custom jobs that can be run on OpenWeights."""
+    mount: Dict[str, str] = {}  # source path -> target path mapping
+    params: Type[BaseModel] = BaseModel  # Pydantic model for parameter validation
+    base_image: str = 'nielsrolf/ow-inference'  # Base Docker image to use
+    requires_vram_gb: int = 24  # Required VRAM in GB
+
+    def __init__(self, client):
+        """Initialize the custom job.
+        `client` should be an instance of `openweights.OpenWeights`."""
+        self.client = client
+
+    @property
+    def _supabase(self):
+        return self.client._supabase
+    
+    @property
+    def _org_id(self):
+        return self.client.organization_id
+
+    def get_entrypoint(self, validated_params: BaseModel) -> str:
+        """Get the entrypoint command for the job.
+        
+        Args:
+            validated_params: The validated parameters as a Pydantic model instance
+        
+        Returns:
+            The command to run as a string
+        """
+        raise NotImplementedError("Subclasses must implement get_entrypoint")
+
+    def _upload_mounted_files(self) -> Dict[str, str]:
+        """Upload all mounted files and return mapping of target paths to file IDs."""
+        uploaded_files = {}
+        
+        for source_path, target_path in self.mount.items():
+            # Handle both files and directories
+            if os.path.isfile(source_path):
+                with open(source_path, 'rb') as f:
+                    file_response = self.client.files.create(f, purpose='custom_job_file')
+                uploaded_files[target_path] = file_response['id']
+            elif os.path.isdir(source_path):
+                # For directories, upload each file maintaining the structure
+                for root, _, files in os.walk(source_path):
+                    for file in files:
+                        full_path = os.path.join(root, file)
+                        rel_path = os.path.relpath(full_path, source_path)
+                        target_file_path = os.path.join(target_path, rel_path)
+                        
+                        with open(full_path, 'rb') as f:
+                            file_response = self.client.files.create(f, purpose='custom_job_file')
+                        uploaded_files[target_file_path] = file_response['id']
+            else:
+                raise ValueError(f"Mount source path does not exist: {source_path}")
+        
+        return uploaded_files
+
+    def create(self, **params) -> Dict[str, Any]:
+        """Create and submit a custom job.
+        
+        Args:
+            **params: Parameters for the job, will be validated against self.params
+
+        Returns:
+            The created job object
+        """
+        # Validate parameters
+        validated_params = self.params(**params)
+        
+        # Upload mounted files
+        mounted_files = self._upload_mounted_files()
+        
+        # Get entrypoint command
+        entrypoint = self.get_entrypoint(validated_params)
+        
+        # Create job
+        job_data = {
+            'type': 'custom',
+            'image': self.base_image,
+            'requires_vram_gb': params.get('requires_vram_gb', self.requires_vram_gb),
+            'script': entrypoint,
+            'params': {
+                'validated_params': validated_params.model_dump(),
+                'mounted_files': mounted_files
+            }
+        }
+        
+        return self.client.jobs.create(**job_data)

openweights/client/events.py

@@ -0,0 +1,42 @@
+from typing import Optional, Dict, Any, List
+
+
+class Events():
+    def __init__(self, supabase):
+        self._supabase = supabase
+    
+    def list(self, job_id: Optional[str]=None, run_id: Optional[str]=None):
+        """List events by job_id or run_id, sorted by created_at in ascending order"""
+        if run_id:
+            query = self._supabase.table('events').select('*').eq('run_id', run_id).order('created_at', desc=False)
+        elif job_id:
+            # First get all runs for this job
+            runs_result = self._supabase.table('runs').select('id').eq('job_id', job_id).execute()
+            run_ids = [run['id'] for run in runs_result.data]
+            # Then get all events for these runs
+            query = self._supabase.table('events').select('*').in_('run_id', run_ids).order('created_at', desc=False)
+        else:
+            raise ValueError("Either job_id or run_id must be provided")
+        
+        result = query.execute()
+        return result.data
+
+    def latest(self, fields: List[str], job_id: Optional[str]=None, run_id: Optional[str]=None) -> Dict[str, Any]:
+        """Get a list of events and return a dict with the latest value for each field"""
+        events = self.list(job_id=job_id, run_id=run_id)
+        latest_values = {}
+        if fields == '*':
+            latest_values = {}
+            for event in events:
+                for key, value in event['data'].items():
+                    if value is not None:
+                        latest_values[key] = value
+        else:
+            events = events[::-1]  # Reverse order to get latest events first
+            while len(fields) > 0 and len(events) > 0:
+                event = events.pop()
+                for field in fields:
+                    if field in event['data']:
+                        latest_values[field] = event['data'][field]
+                        fields.remove(field)
+        return latest_values

openweights/client/files.py

@@ -0,0 +1,106 @@
+from typing import Optional, BinaryIO, Dict, Any, List, Union
+import os
+import hashlib
+from datetime import datetime
+from supabase import Client
+import backoff
+
+from openweights.validate import validate_messages, validate_preference_dataset
+
+
+class Files:
+    def __init__(self, supabase: Client, organization_id: str):
+        self._supabase = supabase
+        self._org_id = organization_id
+
+    def _calculate_file_hash(self, file: BinaryIO) -> str:
+        """Calculate SHA-256 hash of file content"""
+        sha256_hash = hashlib.sha256()
+        for byte_block in iter(lambda: file.read(4096), b""):
+            sha256_hash.update(byte_block)
+        # Add the org ID to the hash to ensure uniqueness
+        sha256_hash.update(self._org_id.encode())
+        file.seek(0)  # Reset file pointer
+        return f"file-{sha256_hash.hexdigest()[:12]}"
+
+    def _get_storage_path(self, file_id: str) -> str:
+        """Get the organization-specific storage path for a file"""
+        try:
+            result = self._supabase.rpc(
+                'get_organization_storage_path',
+                {'org_id': self._org_id, 'filename': file_id}
+            ).execute()
+            return result.data
+        except Exception as e:
+            # Fallback if RPC fails
+            return f"organizations/{self._org_id}/{file_id}"
+
+    @backoff.on_exception(backoff.constant, Exception, interval=1, max_time=60, max_tries=60, on_backoff=lambda details: print(f"Retrying... {details['exception']}"))
+    def create(self, file: BinaryIO, purpose: str) -> Dict[str, Any]:
+        """Upload a file and create a database entry"""
+        file_id = f"{purpose}:{self._calculate_file_hash(file)}"
+
+        # If the file already exists, return the existing file
+        try:
+            existing_file = self._supabase.table('files').select('*').eq('id', file_id).single().execute().data
+            if existing_file:
+                return existing_file
+        except:
+            pass  # File doesn't exist yet, continue with creation
+
+        # Validate file content
+        if not self.validate(file, purpose):
+            raise ValueError("File content is not valid")
+
+        file_size = os.fstat(file.fileno()).st_size
+        filename = getattr(file, 'name', 'unknown')
+
+        # Get organization-specific storage path
+        storage_path = self._get_storage_path(file_id)
+
+        # Store file in Supabase Storage with organization path
+        self._supabase.storage.from_('files').upload(
+            path=storage_path,
+            file=file
+        )
+
+        # Create database entry
+        data = {
+            'id': file_id,
+            'filename': filename,
+            'purpose': purpose,
+            'bytes': file_size
+        }
+        
+        result = self._supabase.table('files').insert(data).execute()
+        
+        return {
+            'id': file_id,
+            'object': 'file',
+            'bytes': file_size,
+            'created_at': datetime.now().timestamp(),
+            'filename': filename,
+            'purpose': purpose,
+        }
+
+    @backoff.on_exception(backoff.constant, Exception, interval=1, max_time=60, max_tries=60, on_backoff=lambda details: print(f"Retrying... {details['exception']}"))
+    def content(self, file_id: str) -> bytes:
+        """Get file content"""
+        storage_path = self._get_storage_path(file_id)
+        return self._supabase.storage.from_('files').download(storage_path)
+    
+    def validate(self, file: BinaryIO, purpose: str) -> bool:
+        """Validate file content"""
+        if purpose in ['conversations']:
+            content = file.read().decode('utf-8')
+            return validate_messages(content)
+        elif purpose == 'preference':
+            content = file.read().decode('utf-8')
+            return validate_preference_dataset(content)
+        else:
+            return True
+        
+    @backoff.on_exception(backoff.constant, Exception, interval=1, max_time=60, max_tries=60, on_backoff=lambda details: print(f"Retrying... {details['exception']}"))
+    def get_by_id(self, file_id: str) -> Dict[str, Any]:
+        """Get file details by ID"""
+        return self._supabase.table('files').select('*').eq('id', file_id).single().execute().data