promote hybrid: 8B map + 30B read + split build (VRAM) + FLUX preload + CUDA cache

arcana/gpu_models.py

@@ -28,21 +28,34 @@ import threading
 
 # fast multi-connection Hub downloads (the GGUFs/FLUX weights crawl otherwise)
 os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
+# Persist the CUDA JIT cache to disk (the Space disk is huge + persists) so a fresh
+# ZeroGPU worker can skip re-JITting llama.cpp's Blackwell kernels — cuts the cold
+# warm-up on later workers.
+os.environ.setdefault("CUDA_CACHE_PATH", "/data/.nv_cache"
+                      if os.path.isdir("/data") else os.path.expanduser("~/.nv_cache"))
+os.environ.setdefault("CUDA_CACHE_MAXSIZE", str(2 * 1024 ** 3))
+try:
+    os.makedirs(os.environ["CUDA_CACHE_PATH"], exist_ok=True)
+except Exception:
+    pass
 
 # ---- model selection (env-overridable) ------------------------------------
+# HYBRID, by necessity: the 30B-Q4-MoE CANNOT reliably emit the dense 22-card×5-field
+# mapping JSON (tested both ways: reasoning = too slow/GPU-aborts; suppressed = leaves
+# fields empty, fails all retries). The dense 8B-Q4 (Llama-arch) fills the JSON
+# reliably and warms fast. So: 8B MAPS, 30B READS (rich prose). Both Nemotron, both
+# llama.cpp (the right engine for the hybrid-Mamba 30B — transformers can't build its
+# Mamba kernels on ZeroGPU). They run in separate @spaces.GPU calls so the 8B+FLUX
+# (build) and the 30B (reading) never share the 48GB GPU.
 _DEF_8B = ("bartowski/nvidia_Llama-3.1-Nemotron-Nano-8B-v1-GGUF",
            "nvidia_Llama-3.1-Nemotron-Nano-8B-v1-Q4_K_M.gguf")
 _DEF_30B = ("unsloth/Nemotron-3-Nano-30B-A3B-GGUF",
             "Nemotron-3-Nano-30B-A3B-Q4_K_M.gguf")
-
-# Mapping model (structured JSON) — the dense 8B (reliable + fast).
-MAP_REPO = os.environ.get("LOCAL_MAP_REPO", os.environ.get("LOCAL_LLM_REPO", _DEF_8B[0]))
-MAP_FILE = os.environ.get("LOCAL_MAP_FILE", os.environ.get("LOCAL_LLM_FILE", _DEF_8B[1]))
+MAP_REPO = os.environ.get("LOCAL_MAP_REPO", _DEF_8B[0])
+MAP_FILE = os.environ.get("LOCAL_MAP_FILE", _DEF_8B[1])
 MAP_N_CTX = int(os.environ.get("LOCAL_MAP_N_CTX", "16384"))
-# Reading model (prose) — defaults to the SAME 8B; set LOCAL_READ_* to the 30B
-# on spaces that should produce the richer readings.
-READ_REPO = os.environ.get("LOCAL_READ_REPO", MAP_REPO)
-READ_FILE = os.environ.get("LOCAL_READ_FILE", MAP_FILE)
+READ_REPO = os.environ.get("LOCAL_READ_REPO", _DEF_30B[0])
+READ_FILE = os.environ.get("LOCAL_READ_FILE", _DEF_30B[1])
 READ_N_CTX = int(os.environ.get("LOCAL_READ_N_CTX", "8192"))
 
 # Distilled klein (built for 4 steps); set LOCAL_IMAGE_MODEL=...-klein-9B for best.
@@ -104,17 +117,50 @@ def _is_nemotron3(fname: str) -> bool:
     return "30b" in f or "nemotron-3-nano-3" in f
 
 
-def _gen(llm, fname: str, system: str, user: str, max_tokens: int, temperature: float) -> str:
-    """Generate clean (non-reasoning) text, picking the prompt style per model."""
-    if _is_nemotron3(fname):
-        # Nemotron-3 ChatML: seed an empty <think></think> to suppress reasoning.
-        prompt = (f"<|im_start|>system\n{system}<|im_end|>\n"
-                  f"<|im_start|>user\n{user}<|im_end|>\n"
-                  f"<|im_start|>assistant\n<think></think>\n")
-        out = llm.create_completion(prompt=prompt, max_tokens=max_tokens,
-                                    temperature=temperature, stop=["<|im_end|>"])
-        return (out["choices"][0]["text"] or "").strip()
-    # Llama-3 template (8B Llama-Nemotron): "detailed thinking off" in system.
+def _last_balanced_json(text: str) -> str:
+    """Return the LAST top-level {...} block in text (the final answer after any
+    reasoning), else the whole text. Lets a reasoning model think first, then we
+    take the JSON it emits at the end."""
+    objs, depth, start, in_str, esc = [], 0, None, False, False
+    for i, c in enumerate(text):
+        if in_str:
+            if esc: esc = False
+            elif c == "\\": esc = True
+            elif c == '"': in_str = False
+            continue
+        if c == '"': in_str = True
+        elif c == "{":
+            if depth == 0: start = i
+            depth += 1
+        elif c == "}":
+            if depth > 0:
+                depth -= 1
+                if depth == 0 and start is not None:
+                    objs.append(text[start:i + 1]); start = None
+    return objs[-1] if objs else text
+
+
+def _suppress_gen(llm, system: str, user: str, max_tokens: int, temperature: float) -> str:
+    """Nemotron-3 ChatML with an empty <think></think> seed → clean, non-reasoning
+    prose (fast). Used for readings."""
+    prompt = (f"<|im_start|>system\n{system}<|im_end|>\n"
+              f"<|im_start|>user\n{user}<|im_end|>\n"
+              f"<|im_start|>assistant\n<think></think>\n")
+    out = llm.create_completion(prompt=prompt, max_tokens=max_tokens,
+                                temperature=temperature, stop=["<|im_end|>"])
+    return (out["choices"][0]["text"] or "").strip()
+
+
+def llama_map(system: str, user: str, max_tokens: int = 12000,
+              temperature: float = 0.5) -> str:
+    """Structured JSON mapping — think-SUPPRESSED (fast). Reasoning-then-extract was
+    far too slow (16k tokens × 2 passes → GPU-time aborts). Suppressed generation
+    with a generous token budget should be clean: the earlier 'empty concept'
+    failures were almost certainly JSON truncation from too-low max_tokens, not the
+    model. designer.extract_json + a couple retries cover any residual slip."""
+    llm = _get_llama(MAP_REPO, MAP_FILE, MAP_N_CTX)
+    if _is_nemotron3(MAP_FILE):
+        return _suppress_gen(llm, system, user, max_tokens, temperature)
     out = llm.create_chat_completion(
         messages=[{"role": "system", "content": "detailed thinking off\n\n" + system},
                   {"role": "user", "content": user}],
@@ -122,19 +168,17 @@ def _gen(llm, fname: str, system: str, user: str, max_tokens: int, temperature:
     return (out["choices"][0]["message"]["content"] or "").strip()
 
 
-def llama_map(system: str, user: str, max_tokens: int = 10000,
-              temperature: float = 0.4) -> str:
-    """Structured JSON mapping via the small dense model (clean JSON, no grammar —
-    the dense 8B emits valid JSON and designer.extract_json + retries cover slips)."""
-    llm = _get_llama(MAP_REPO, MAP_FILE, MAP_N_CTX)
-    return _gen(llm, MAP_FILE, system, user, max_tokens, temperature)
-
-
 def llama_read(system: str, user: str, max_tokens: int = 9000,
                temperature: float = 0.7) -> str:
-    """Prose readings via the reading model (the 30B where configured → richer)."""
+    """Prose readings (think-suppressed for clean, fast output)."""
     llm = _get_llama(READ_REPO, READ_FILE, READ_N_CTX)
-    return _gen(llm, READ_FILE, system, user, max_tokens, temperature)
+    if _is_nemotron3(READ_FILE):
+        return _suppress_gen(llm, system, user, max_tokens, temperature)
+    out = llm.create_chat_completion(
+        messages=[{"role": "system", "content": "detailed thinking off\n\n" + system},
+                  {"role": "user", "content": user}],
+        max_tokens=max_tokens, temperature=temperature)
+    return (out["choices"][0]["message"]["content"] or "").strip()
 
 
 def preload_pipe():