import time import logging from enum import Enum from typing import List, Optional class NvidiaModel(str, Enum): # Primary Heavy Coding # # glm-5.1 was RETIRED upstream — NVIDIA answers HTTP 410 for it on every key, # which silently broke the entire Gemini->NVIDIA fallback chain. Verified live # on all 15 keys. glm-5.2 is the current model and returns correct answers, # but it is a reasoning model and measured ~270s for a trivial prompt, so it # must NOT be used on an interactive path (HF's gateway cuts at ~150s). GLM_5_2 = "z-ai/glm-5.2" # Kept so any existing reference to the old name keeps working; it now points # at the live model rather than the retired one. GLM_5_1 = "z-ai/glm-5.2" MINIMAX_M3 = "minimaxai/minimax-m3" MINIMAX_M2_7 = "minimaxai/minimax-m2.7" DEEPSEEK_V4_PRO = "deepseek-ai/deepseek-v4-pro" DEEPSEEK_V4_FLASH = "deepseek-ai/deepseek-v4-flash" NEMOTRON_3_ULTRA_550B = "nvidia/nemotron-3-ultra-550b-a55b" # Secondary Agentic NEMOTRON_3_SUPER_120B = "nvidia/nemotron-3-super-120b-a12b" QWEN_3_5_122B = "qwen/qwen3.5-122b-a10b" KIMI_K2_6 = "moonshotai/kimi-k2.6" MISTRAL_MEDIUM_3_5_128B = "mistralai/mistral-medium-3.5-128b" MISTRAL_SMALL_4_119B = "mistralai/mistral-small-4-119b-2603" GEMMA_4_31B = "google/gemma-4-31b-it" STEP_3_7_FLASH = "stepfun-ai/step-3.7-flash" # Multimodal / Vision / Audio NEMOTRON_NANO_OMNI = "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning" COSMOS3_NANO_REASONER = "nvidia/cosmos-reason2-8b" COSMOS3_NANO = "microsoft/phi-3-vision-128k-instruct" NEMOTRON_VOICECHAT = "nvidia/nemotron-nano-12b-v2-vl" # Specialized NEMOTRON_3_5_CONTENT_SAFETY = "nvidia/nemotron-3.5-content-safety" NEMOTRON_3_CONTENT_SAFETY = "nvidia/nemotron-3-content-safety" GLINER_PII = "nvidia/gliner-pii" SYNTHETIC_VIDEO_DETECTOR = "nvidia/ai-synthetic-video-detector" ACTIVE_SPEAKER_DETECTION = "nvidia/nemotron-nano-12b-v2-vl" ISING_CALIBRATION_35B = "nvidia/nemotron-nano-12b-v2-vl" class KeyTier(str, Enum): HEAVY_COMPUTE = "heavy_compute" # For GLM, DeepSeek, Nemotron Ultra AGENTIC = "agentic" # For Qwen, Kimi, Mistral MULTIMODAL = "multimodal" # For Cosmos, Omni, VoiceChat SPECIALIZED = "specialized" # For Safety, PII, Detector MAX_NVIDIA_KEYS = 15 # NVIDIA NIM specific rate limits TPM_LIMIT = 50000 # Tokens Per Minute RPM_LIMIT = 50 # Requests Per Minute class NvidiaKeyTracker: def __init__(self, key: str, tier: KeyTier): self.key = key self.tier = tier self.status = "healthy" self.tokens_used_this_minute = 0 self.requests_this_minute = 0 self.minute_start_time = time.time() self.last_fail_time = 0 self.fail_reason = "" def reset_limits_if_needed(self): if time.time() - self.minute_start_time > 60: self.tokens_used_this_minute = 0 self.requests_this_minute = 0 self.minute_start_time = time.time() def can_accommodate(self, estimated_tokens: int = 1000) -> bool: self.reset_limits_if_needed() if self.status != "healthy": return False if self.requests_this_minute >= RPM_LIMIT: return False if self.tokens_used_this_minute + estimated_tokens > TPM_LIMIT: return False return True class NvidiaKeyManager: """Advanced manager mapping specific keys to specific model tiers to prevent heavy models from throttling specialized models.""" def __init__(self): self.key_trackers: List[NvidiaKeyTracker] = [] self._load_keys() def _load_keys(self): try: import json, os from backend.services.usb_monitor import get_db_path # Determine runtime location to partition keys correctly is_cloud = bool(os.environ.get("SPACE_ID")) suffix = "CLOUD" if is_cloud else "PC" # First try the JSON file (local PC) path = os.path.join(os.path.dirname(get_db_path()), "nvidia_keys.json") if os.path.exists(path): with open(path, "r") as f: keys = json.load(f) # If JSON has exactly 14 keys, split them half-half if len(keys) >= 14: if is_cloud: keys = keys[7:14] else: keys = keys[0:7] for k in keys: self.inject_key(k) # Then layer in HF Secrets env vars. # We look for NVIDIA_API_KEY_1_PC/CLOUD up to 7, mirroring Google architecture. for i in range(1, 8): # Try new partitioned format first (e.g. NVIDIA_API_KEY_1_CLOUD) env_key = os.environ.get(f"NVIDIA_API_KEY_{i}_{suffix}", "").strip() # If they haven't renamed them yet, gracefully fallback to splitting the 1..15 raw keys if not env_key: # Cloud gets 8-14, PC gets 1-7 raw_index = i + 7 if is_cloud else i env_key = os.environ.get(f"NVIDIA_API_KEY_{raw_index}", "").strip() if env_key: self.inject_key(env_key) # Sweep any remaining numbered keys the partitioned ranges above don't cover. # The loop reads 1-7 on PC and 8-14 on cloud, so NVIDIA_API_KEY_15 was # configured (confirmed present live via /system/key_diagnostics) but never # actually loaded by either branch — a paid key sitting unused. Rather than # widen one range and risk the PC/cloud split overlapping, pick up anything # left over here; inject_key() already de-duplicates and caps at 15, so a key # the partitioned loop already loaded is ignored harmlessly. for i in range(1, 16): leftover = os.environ.get(f"NVIDIA_API_KEY_{i}", "").strip() if leftover: self.inject_key(leftover) # Single fallback key single = os.environ.get("NVIDIA_API_KEY", "").strip() if single: self.inject_key(single) except Exception as e: logging.error(f"NvidiaKeyManager: Failed to load keys: {e}") def _save_keys(self): try: import json import os from backend.services.usb_monitor import get_db_path path = os.path.join(os.path.dirname(get_db_path()), "nvidia_keys.json") with open(path, "w") as f: json.dump([t.key for t in self.key_trackers], f) except Exception as e: logging.error(f"NvidiaKeyManager: Failed to save keys: {e}") def inject_key(self, key: str): if any(t.key == key for t in self.key_trackers): return current_count = len(self.key_trackers) if current_count >= 15: logging.warning("NvidiaKeyManager: Max 15 keys reached. Ignoring injection.") return # Round-robin distribution ensures equal distribution of keys across all tiers # regardless of whether the user provided 7 keys (partitioned) or 15 keys (raw) tier_map = [ KeyTier.HEAVY_COMPUTE, KeyTier.AGENTIC, KeyTier.MULTIMODAL, KeyTier.SPECIALIZED ] tier = tier_map[current_count % 4] tracker = NvidiaKeyTracker(key, tier) self.key_trackers.append(tracker) self._save_keys() logging.info(f"NvidiaKeyManager: Injected key {current_count + 1} into tier {tier.value}.") def _get_tier_for_model(self, model: NvidiaModel) -> KeyTier: heavy_models = [NvidiaModel.GLM_5_2, NvidiaModel.MINIMAX_M3, NvidiaModel.MINIMAX_M2_7, NvidiaModel.DEEPSEEK_V4_PRO, NvidiaModel.DEEPSEEK_V4_FLASH, NvidiaModel.NEMOTRON_3_ULTRA_550B] agentic_models = [NvidiaModel.NEMOTRON_3_SUPER_120B, NvidiaModel.QWEN_3_5_122B, NvidiaModel.KIMI_K2_6, NvidiaModel.MISTRAL_MEDIUM_3_5_128B, NvidiaModel.MISTRAL_SMALL_4_119B, NvidiaModel.GEMMA_4_31B, NvidiaModel.STEP_3_7_FLASH] multimodal_models = [NvidiaModel.NEMOTRON_NANO_OMNI, NvidiaModel.COSMOS3_NANO_REASONER, NvidiaModel.COSMOS3_NANO, NvidiaModel.NEMOTRON_VOICECHAT] if model in heavy_models: return KeyTier.HEAVY_COMPUTE if model in agentic_models: return KeyTier.AGENTIC if model in multimodal_models: return KeyTier.MULTIMODAL return KeyTier.SPECIALIZED def get_optimal_key_for_model(self, model: NvidiaModel, estimated_tokens: int = 1000) -> Optional[str]: target_tier = self._get_tier_for_model(model) # Phase 1: Try to find a healthy key in the correct tier that respects TPM/RPM limits for tracker in self.key_trackers: if tracker.tier == target_tier and tracker.can_accommodate(estimated_tokens): return tracker.key # Phase 2: Automatic Failover across tiers if the native tier is fully exhausted # We allow Agentic to fallback to Heavy Compute keys if needed. for tracker in self.key_trackers: if tracker.can_accommodate(estimated_tokens): logging.warning(f"NvidiaKeyManager: {target_tier.value} keys exhausted. Failing over to {tracker.tier.value} key.") return tracker.key # Phase 3: Attempt Auto-Recovery of 429 failed keys for tracker in self.key_trackers: if tracker.status == "failed" and tracker.fail_reason == "429": if time.time() - tracker.last_fail_time > 60: # Recover after 1 min for rate limits tracker.status = "healthy" logging.info("NvidiaKeyManager: Auto-recovered key from 429 status.") return tracker.key # S4 Phase 4: no permanent key death. Other 4xx failures (transient auth # blips, gateway hiccups mislabelled as 4xx) get a 10-minute cooldown # instead of being lost until process restart. for tracker in self.key_trackers: if tracker.status == "failed" and time.time() - tracker.last_fail_time > 600: tracker.status = "healthy" logging.info(f"NvidiaKeyManager: Auto-recovered key after cooldown (was {tracker.fail_reason}).") return tracker.key return None def mark_key_failed(self, key: str, status_code: int): for tracker in self.key_trackers: if tracker.key == key: tracker.status = "failed" tracker.last_fail_time = time.time() tracker.fail_reason = str(status_code) logging.warning(f"NvidiaKeyManager: Key in {tracker.tier.value} failed with {status_code}.") break def record_usage(self, key: str, tokens: int): for tracker in self.key_trackers: if tracker.key == key: tracker.tokens_used_this_minute += tokens tracker.requests_this_minute += 1 break nvidia_key_manager = NvidiaKeyManager() def call_nvidia_model(prompt: str, initial_model: NvidiaModel = NvidiaModel.GLM_5_2) -> str: from openai import OpenAI import httpx # Dynamically determine fallback models based on the initial model's tier target_tier = nvidia_key_manager._get_tier_for_model(initial_model) # We reconstruct the tier lists here to use them for dynamic fallback heavy_models = [NvidiaModel.GLM_5_2, NvidiaModel.MINIMAX_M3, NvidiaModel.MINIMAX_M2_7, NvidiaModel.DEEPSEEK_V4_PRO, NvidiaModel.DEEPSEEK_V4_FLASH, NvidiaModel.NEMOTRON_3_ULTRA_550B] agentic_models = [NvidiaModel.NEMOTRON_3_SUPER_120B, NvidiaModel.QWEN_3_5_122B, NvidiaModel.KIMI_K2_6, NvidiaModel.MISTRAL_MEDIUM_3_5_128B, NvidiaModel.MISTRAL_SMALL_4_119B, NvidiaModel.GEMMA_4_31B, NvidiaModel.STEP_3_7_FLASH] multimodal_models = [NvidiaModel.NEMOTRON_NANO_OMNI, NvidiaModel.COSMOS3_NANO_REASONER, NvidiaModel.COSMOS3_NANO, NvidiaModel.NEMOTRON_VOICECHAT] if target_tier == KeyTier.HEAVY_COMPUTE: tier_models = heavy_models elif target_tier == KeyTier.AGENTIC: tier_models = agentic_models elif target_tier == KeyTier.MULTIMODAL: tier_models = multimodal_models else: tier_models = [initial_model] # Specialized tier has no explicit failover list defined here fallback_models = [initial_model] + [m for m in tier_models if m != initial_model] # Remove duplicates while preserving order models_to_try = [] for m in fallback_models: if m not in models_to_try: models_to_try.append(m) last_error = "" for current_model in models_to_try: # For each model, we allow up to 3 key-failovers (in case keys are hitting 429s) for attempt in range(3): key = nvidia_key_manager.get_optimal_key_for_model(current_model) if not key: last_error = "No healthy keys available across any tier." break # Break inner loop, try next model (though if no keys exist, it likely won't help) try: client = OpenAI( base_url="https://integrate.api.nvidia.com/v1", api_key=key, http_client=httpx.Client(timeout=30.0) ) response = client.chat.completions.create( model=current_model.value, messages=[{"role": "user", "content": prompt}], temperature=0.2, # S4: reasoning models (glm-5.1 etc.) spend their budget on # think-tokens before the answer — 1024 truncated structured # JSON outputs (Builder blueprints) mid-object every time. max_tokens=8192 ) total_tokens = response.usage.total_tokens if response.usage else 1000 nvidia_key_manager.record_usage(key, total_tokens) return response.choices[0].message.content except Exception as e: status_code = getattr(e, 'status_code', 500) if 'Timeout' in str(type(e)): status_code = 504 last_error = str(e) # S4: 404/410 mean the MODEL is gone/retired — the key is fine. # Marking keys failed here let one dead model poison the whole # fleet ("No healthy keys available across any tier"). Rotate model. if status_code in (404, 410): logging.warning(f"NvidiaKeyManager: Model {current_model.value} returned {status_code} (model unavailable). Failing over to next model...") break # Break inner loop, move to next model in outer loop # If Auth/Rate Limit (4xx), the KEY is burnt. Mark key failed and retry same model with new key. if 400 <= status_code < 500: logging.warning(f"NvidiaKeyManager: Key failed on {current_model.value} with {status_code}. Marking key failed and rotating...") nvidia_key_manager.mark_key_failed(key, status_code) continue # Next attempt in inner loop (same model, new key) # If Server Error/Timeout (5xx), the MODEL is down. Break inner loop and try NEXT model. else: logging.warning(f"NvidiaKeyManager: Model {current_model.value} returned {status_code}. Server/Timeout issue. Failing over to next model...") break # Break inner loop, move to next model in outer loop return f"[NVIDIA FALLBACK FAILED] Exhausted all failover models and keys. Last error: {last_error}"