| """Qwen transport with bounded concurrency and conservative spend reservations.""" |
|
|
| import base64 |
| import io |
| import json |
| import threading |
| import time |
|
|
| import requests |
| from PIL import Image |
|
|
| from lerobot.annotations.steerable_pipeline.config import VlmConfig |
| from lerobot.annotations.steerable_pipeline.vlm_client import _GenericTextClient, _to_openai_messages |
|
|
|
|
| class BudgetExceededError(RuntimeError): |
| pass |
|
|
|
|
| class DeadlineExceededError(BudgetExceededError): |
| pass |
|
|
|
|
| class Budget: |
| def __init__( |
| self, |
| limit, |
| input_rate, |
| output_rate, |
| cached_input_rate=None, |
| context_length=262144, |
| pricing_tiers=None, |
| ): |
| self.limit, self.input_rate, self.output_rate = limit, input_rate, output_rate |
| self.reserved, self.measured, self.calls = 0.0, 0.0, 0 |
| self.cached_input_rate = input_rate if cached_input_rate is None else cached_input_rate |
| self.context_length = context_length |
| self.pricing_tiers = sorted(pricing_tiers or [], key=lambda t: t["min_input_tokens"]) |
| self.input_tokens = self.output_tokens = self.cached_tokens = self.attempts = self.rate_limits = 0 |
| self.media_attempts = self.media_completed = 0 |
| self.media_usd = 0.0 |
| self.lock = threading.Lock() |
|
|
| def reserve(self, max_tokens): |
| |
| amount = self.reservation_cost(max_tokens) |
| with self.lock: |
| if self.reserved + amount > self.limit: |
| raise BudgetExceededError("Spending target reached; completed results are saved.") |
| self.reserved += amount |
| self.attempts += 1 |
| return amount |
|
|
| def reservation_cost(self, max_tokens): |
| input_rate = max([self.input_rate] + [t["input_rate"] for t in self.pricing_tiers]) |
| output_rate = max([self.output_rate] + [t["output_rate"] for t in self.pricing_tiers]) |
| return (self.context_length * input_rate + max_tokens * output_rate) / 1e6 |
|
|
| def reserve_media(self, amount): |
| """Reserve a fixed-price media request (image edit or video). Never refunded.""" |
| with self.lock: |
| if self.reserved + amount > self.limit: |
| raise BudgetExceededError("Spending target reached; completed results are saved.") |
| self.reserved += amount |
| self.media_attempts += 1 |
| return amount |
|
|
| def record_media(self, reservation, cost): |
| """A media request returned its output: replace its reservation with its list price.""" |
| with self.lock: |
| self.reserved += cost - reservation |
| self.measured += cost |
| self.media_usd += cost |
| self.media_completed += 1 |
|
|
| def reconcile(self, reservation, usage): |
| if not isinstance(usage, dict): |
| return |
| values = [usage.get("prompt_tokens"), usage.get("completion_tokens")] |
| if any(type(v) is not int or v < 0 for v in values): |
| return |
| cached = usage.get( |
| "cached_tokens", (usage.get("prompt_tokens_details") or {}).get("cached_tokens", 0) |
| ) |
| cached = min(values[0], max(0, cached)) if type(cached) is int else 0 |
| input_rate, output_rate, cached_rate = self.input_rate, self.output_rate, self.cached_input_rate |
| for tier in self.pricing_tiers: |
| if values[0] >= tier["min_input_tokens"]: |
| input_rate, output_rate, cached_rate = ( |
| tier["input_rate"], |
| tier["output_rate"], |
| tier["cached_input_rate"], |
| ) |
| cost = ((values[0] - cached) * input_rate + cached * cached_rate + values[1] * output_rate) / 1e6 |
| with self.lock: |
| self.reserved += cost - reservation |
| self.measured += cost |
| self.calls += 1 |
| self.input_tokens += values[0] |
| self.output_tokens += values[1] |
| self.cached_tokens += cached |
|
|
| def summary(self): |
| with self.lock: |
| return { |
| "measured_inference_usd": self.measured, |
| "reserved_including_uncertain_usd": self.reserved, |
| "successful_calls": self.calls, |
| "inference_target_usd": self.limit, |
| "input_tokens": self.input_tokens, |
| "output_tokens": self.output_tokens, |
| "cached_input_tokens": self.cached_tokens, |
| "attempted_requests": self.attempts, |
| "requests_without_usage": self.attempts - self.calls, |
| "rate_limited_requests": self.rate_limits, |
| "media_requests": self.media_attempts, |
| "completed_media_requests": self.media_completed, |
| "media_usd": self.media_usd, |
| } |
|
|
|
|
| def compact_images(messages, max_bytes=2_800_000): |
| if len(json.dumps(messages).encode()) <= max_bytes: |
| return messages |
| images = [] |
| for message in messages: |
| if not isinstance(message.get("content"), list): |
| continue |
| for part in message["content"]: |
| if part.get("type") == "image_url": |
| url = part["image_url"]["url"] |
| if not url.startswith("data:image/"): |
| raise ValueError("Only decoded dataset images can be submitted.") |
| images.append( |
| (part, Image.open(io.BytesIO(base64.b64decode(url.split(",", 1)[1]))).convert("RGB")) |
| ) |
| for quality in (84, 65, 45, 30): |
| for part, image in images: |
| buffer = io.BytesIO() |
| image.save(buffer, format="JPEG", quality=quality, optimize=True) |
| part["image_url"]["url"] = ( |
| "data:image/jpeg;base64," + base64.b64encode(buffer.getvalue()).decode() |
| ) |
| if len(json.dumps(messages).encode()) <= max_bytes: |
| return messages |
| raise ValueError("Contact sheets exceed the provider request limit.") |
|
|
|
|
| class Provider: |
| def __init__(self, token, model, budget, deadline=None, concurrency=2, generation_options=None): |
| self.token, self.model, self.budget = token, model, budget |
| |
| capacity = max(1, int(budget.limit / budget.reservation_cost(4096))) |
| self.concurrency = max(1, min(int(concurrency), capacity, 32)) |
| self.semaphore = threading.BoundedSemaphore(self.concurrency) |
| self.deadline = deadline |
| |
| |
| self.generation_options = ( |
| {"chat_template_kwargs": {"enable_thinking": False}, "seed": 1729} |
| if generation_options is None and model.startswith("Qwen/") |
| else dict(generation_options or {}) |
| ) |
| self.client = _GenericTextClient(self.generate, VlmConfig(max_new_tokens=4096, temperature=0.2)) |
|
|
| def generate_json(self, *args, **kwargs): |
| return self.client.generate_json(*args, **kwargs) |
|
|
| def remaining_time(self): |
| left = self.deadline - time.time() if self.deadline else 300 |
| if left <= 1: |
| raise DeadlineExceededError("Runtime limit reached; saving completed episodes.") |
| return left |
|
|
| def generate(self, batch, max_tokens, temperature): |
| output = [] |
| for messages in batch: |
| self.remaining_time() |
| content, _ = _to_openai_messages(messages) |
| body = { |
| "model": self.model, |
| "messages": compact_images(content), |
| "max_tokens": max_tokens, |
| "temperature": temperature, |
| **self.generation_options, |
| } |
| with self.semaphore: |
| for attempt in range(8): |
| left = self.remaining_time() |
| reservation = self.budget.reserve(max_tokens) |
| try: |
| response = requests.post( |
| "https://router.huggingface.co/v1/chat/completions", |
| headers={"Authorization": "Bearer " + self.token}, |
| json=body, |
| timeout=(min(20, left), min(300, left)), |
| ) |
| except requests.RequestException: |
| self.remaining_time() |
| |
| raise RuntimeError( |
| "Inference connection failed; charges are reserved conservatively." |
| ) from None |
| if response.status_code == 429: |
| |
| with self.budget.lock: |
| self.budget.rate_limits += 1 |
| time.sleep(min(2**attempt, 30, self.remaining_time())) |
| continue |
| if response.status_code >= 400: |
| raise RuntimeError(f"Inference provider returned HTTP {response.status_code}.") |
| data = response.json() |
| self.budget.reconcile(reservation, data.get("usage")) |
| choice = data["choices"][0] |
| if choice.get("finish_reason") == "length": |
| raise ValueError("The model response exceeded its output limit.") |
| output.append(choice["message"].get("content") or "") |
| break |
| else: |
| raise RuntimeError("Inference provider remained rate limited.") |
| return output |
|
|
