"""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): # Reserve the model's entire context, not an optimistic image-token guess. 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 # Unknown charges keep the complete reservation. 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 # Avoid exhausting a small budget merely by reserving many in-flight calls. 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 # Older manifests used Qwen-specific parameters. New jobs pin each route's # tested parameters; an explicit empty dict means no extra parameters. 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() # The provider may have processed the request. Never release its reservation. raise RuntimeError( "Inference connection failed; charges are reserved conservatively." ) from None if response.status_code == 429: # Rejections may be ambiguous; retain the reservation before retrying. 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