"""Generation backends for the NEMOCITY mind.
Protocol (duck-typed; see `Backend`):
name: str # "mock" | "llamacpp" | "zerogpu"
model_id: str # human-readable model identifier
oneshot: bool # always True — NEMOCITY is one-shot only
async def generate_stream(prompt: str, grammar: str | None,
max_tokens: int) -> AsyncIterator[str]
`grammar` is a plain string TAG (mind/validate.py), never a real grammar:
"build" | "fix" | "moderation" select the wanted JSON shape (validation +
temperature 0.4), a "!strict" suffix means a retry pass at temperature 0,
None means free text. The protocol signature is byte-compatible with
godseed's so the copied plumbing keeps working.
Selected via the CITY_BACKEND env: mock (default) | zerogpu | llamacpp.
- MockBackend: deterministic keyword-driven scripts (mind/mock_scripts.py)
with small delays so streaming looks alive. Local dev + headless demos.
- ZeroGPUBackend: transformers + spaces.GPU, Nemotron-Nano-9B-v2 resident on
cuda. ONE @spaces.GPU call per petition (the June 12 lesson). Ship path.
- LlamaCppBackend: llama-cpp-python + a local GGUF; offgrid stretch, not the
ship path.
"""
from __future__ import annotations
import asyncio
import concurrent.futures
import json
import os
import re
import threading
from typing import AsyncIterator, Protocol, runtime_checkable
from . import prompts
from .mock_scripts import build_script, fix_script
from .validate import (
FIX_TAG,
base_tag,
is_moderation_grammar,
is_strict,
parse_build,
parse_fix,
parse_moderation,
)
@runtime_checkable
class Backend(Protocol):
"""Structural type for generation backends."""
name: str
model_id: str
def generate_stream(
self, prompt: str, grammar: str | None, max_tokens: int
) -> AsyncIterator[str]: ...
# --------------------------------------------------------------------------
# Mock backend — what demos run on
# --------------------------------------------------------------------------
_MOCK_DENYLIST: tuple[tuple[re.Pattern[str], str], ...] = (
(re.compile(r"\b(nazi|hitler|swastika|kkk|genocide|lynch)\b", re.I), "hate"),
(re.compile(r"\b(kill|murder|slaughter|massacre|behead|stab)\b", re.I), "violence"),
(re.compile(r"\b(bomb|terror|terrorist|terrorism|shooting)\b", re.I), "violence"),
(re.compile(r"\b(rape|porn|sex|sexual|nude|naked|nsfw)\b", re.I), "sexual"),
(re.compile(r"\b(suicide|self[- ]?harm|cutting)\b", re.I), "self-harm"),
)
def _chunks(text: str, size: int) -> list[str]:
return [text[i : i + size] for i in range(0, len(text), size)]
class MockBackend:
"""Deterministic scripted backend.
Routing (marker constants from mind/prompts.py):
- moderation tag -> keyword-denylist verdict on the text after
CANDIDATE_MARKER;
- prompts containing FIX_MARKER -> fix script (reads the rendered
candidates/numbers back out of the prompt);
- everything else -> BUILD JSON keyed on the petition line after the
last PETITION_MARKER.
Token delays come from CITY_MOCK_DELAY (seconds per chunk, default
0.012); tests pass delay=0.0 explicitly.
"""
name = "mock"
model_id = "nemocity-mock-scripts"
oneshot = True
def __init__(self, delay: float | None = None):
if delay is None:
try:
delay = float(os.environ.get("CITY_MOCK_DELAY", "0.012"))
except ValueError:
delay = 0.012
self.delay = max(0.0, delay)
async def _tick(self) -> None:
if self.delay:
await asyncio.sleep(self.delay)
else:
await asyncio.sleep(0) # stay fair to the event loop
@staticmethod
def _tail(prompt: str, marker: str) -> str:
parts = prompt.rsplit(marker, 1)
return parts[-1] if len(parts) == 2 else prompt
def _petition(self, prompt: str) -> str:
tail = self._tail(prompt, prompts.PETITION_MARKER)
return (tail.splitlines() or [""])[0].strip()
def _moderate(self, prompt: str) -> dict:
tail = self._tail(prompt, prompts.CANDIDATE_MARKER)
candidate = (tail.splitlines() or [""])[0].strip()
for pattern, category in _MOCK_DENYLIST:
if pattern.search(candidate):
return {"allowed": False, "category": category}
return {"allowed": True, "category": ""}
async def generate_stream(
self, prompt: str, grammar: str | None = None, max_tokens: int = 700
) -> AsyncIterator[str]:
if is_moderation_grammar(grammar):
verdict = json.dumps(self._moderate(prompt))
for chunk in _chunks(verdict, 8):
await self._tick()
yield chunk
return
if prompts.FIX_MARKER in prompt:
payload = fix_script(prompt)
else:
payload = build_script(self._petition(prompt))
for chunk in _chunks(payload, 12):
await self._tick()
yield chunk
# --------------------------------------------------------------------------
# llama.cpp backend — local GGUF (offgrid stretch; NOT the ship path)
# --------------------------------------------------------------------------
_THREAD_DONE = object()
class LlamaCppBackend:
"""llama-cpp-python backend (CPU). No grammars — tags only pick the
temperature; the lenient parsers + planner fallback cover messy output.
Lazy: nothing is imported or downloaded until the first generate call.
Model resolution order:
1. CITY_GGUF env — local path to a .gguf file;
2. hf_hub_download(CITY_GGUF_REPO, CITY_GGUF_FILE).
Generation runs in a thread executor; chunks cross into the event loop
through a bounded asyncio.Queue (backpressure included). If the consumer
abandons the stream early, a cancel flag stops the producer thread.
"""
name = "llamacpp"
oneshot = True
DEFAULT_REPO = "nvidia/NVIDIA-Nemotron-3-Nano-4B-GGUF"
DEFAULT_FILE = "NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf"
def __init__(
self,
model_path: str | None = None,
n_ctx: int = 4096,
n_threads: int | None = None,
):
self.model_path = model_path or os.environ.get("CITY_GGUF") or None
self.repo_id = os.environ.get("CITY_GGUF_REPO", self.DEFAULT_REPO)
self.filename = os.environ.get("CITY_GGUF_FILE", self.DEFAULT_FILE)
self.n_ctx = n_ctx
self.n_threads = n_threads
self.model_id = self.model_path or f"{self.repo_id}:{self.filename}"
self._llm = None
self._lock = threading.Lock()
# -- lazy model load (thread-safe; called from executor threads) --------
def _ensure_llm(self):
with self._lock:
if self._llm is not None:
return self._llm
from llama_cpp import Llama # heavy import deferred
path = self.model_path
if not path or not os.path.exists(path):
from huggingface_hub import hf_hub_download
path = hf_hub_download(repo_id=self.repo_id, filename=self.filename)
kwargs = {"model_path": path, "n_ctx": self.n_ctx, "verbose": False}
threads = self.n_threads or (os.cpu_count() or 4)
kwargs["n_threads"] = threads
kwargs["n_batch"] = 512 # larger prompt-eval batch -> faster prefill
self._llm = Llama(**kwargs)
self.model_id = path
return self._llm
async def generate_stream(
self, prompt: str, grammar: str | None = None, max_tokens: int = 700
) -> AsyncIterator[str]:
loop = asyncio.get_running_loop()
queue: asyncio.Queue = asyncio.Queue(maxsize=64)
cancelled = threading.Event()
if grammar is None:
temperature = 0.7
else:
temperature = 0.0 if is_strict(grammar) else 0.35
def _put(item) -> bool:
"""Push from the producer thread; honor cancellation."""
future = asyncio.run_coroutine_threadsafe(queue.put(item), loop)
while not cancelled.is_set():
try:
future.result(timeout=0.25)
return True
except concurrent.futures.TimeoutError:
continue
except Exception:
return False
future.cancel()
return False
def _produce() -> None:
try:
llm = self._ensure_llm()
stream = llm.create_completion(
prompt=prompt,
max_tokens=max_tokens,
stream=True,
temperature=temperature,
top_p=0.9,
repeat_penalty=1.1,
)
for part in stream:
if cancelled.is_set():
break
text = part["choices"][0].get("text", "")
if text and not _put(text):
break
except Exception as exc: # surfaced to the consumer below
_put(exc)
finally:
_put(_THREAD_DONE)
producer = loop.run_in_executor(None, _produce)
try:
while True:
item = await queue.get()
if item is _THREAD_DONE:
break
if isinstance(item, Exception):
raise item
yield item
finally:
cancelled.set()
# Drain so a blocked producer can reach its sentinel and exit.
while not queue.empty():
queue.get_nowait()
await producer
# --------------------------------------------------------------------------
# ZeroGPU backend — Nemotron-Nano-9B-v2 (the ship path)
# --------------------------------------------------------------------------
class ZeroGPUBackend:
"""transformers + spaces.GPU backend (HF ZeroGPU).
All heavy imports (torch, transformers, spaces) are deferred to the first
generate call so this module imports cleanly anywhere. The GPU-decorated
function is built lazily; if `spaces` is unavailable (local dev) the bare
function is used.
No grammar support, so when a tag is supplied the output is validated
(mind/validate.py) and regenerated once at temperature ~0 on failure. If
both attempts are malformed the raw text is yielded anyway — the
planner's retry/fallback path handles it from there.
Streaming note: tokens are buffered per generation (validation requires
the full text) and then yielded in small chunks, so SSE consumers still
see incremental output.
"""
name = "zerogpu"
oneshot = True # a petition is ONE @spaces.GPU call (ZeroGPU detaches the
# GPU between calls; godseed's per-turn loop of ~8 calls crashed on #2)
# Nemotron-Nano-9B-v2: its nemotron_h arch is IN-TREE in transformers, so it
# loads with trust_remote_code=False and runs the kernel-free native Mamba-2
# path. The Nemotron-3 4B/30B repos hard-require remote code + mamba-ssm
# kernels, which cannot import on the ZeroGPU image (June 12).
DEFAULT_MODEL = "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
GPU_DURATION_S = 240
def __init__(self, model_id: str | None = None, max_input_tokens: int = 4096):
self.model_id = model_id or os.environ.get(
"CITY_HF_MODEL", self.DEFAULT_MODEL
)
self.max_input_tokens = max_input_tokens
self._tokenizer = None
self._model = None
self._gpu_generate = None
self._on_cuda = False
self._lock = threading.Lock()
# Load weights to CPU at construction (app startup). CRITICAL: do NOT
# touch CUDA here. Initializing CUDA in the main process poisons
# ZeroGPU's per-request GPU attach (the forked worker inherits a broken
# context -> "NVML_SUCCESS INTERNAL ASSERT FAILED" in the CUDA allocator
# on every forward, June 12). The model moves to cuda guarded by
# `import spaces` below, which virtualizes the placement.
self._ensure()
# -- setup (model resident; cuda placement guarded by spaces) -----------
def _ensure(self):
with self._lock:
if self._gpu_generate is not None:
return self._gpu_generate
import torch # noqa: F401
from transformers import AutoModelForCausalLM, AutoTokenizer
# trust_remote_code MUST stay False: the repo also ships custom code
# that hard-requires mamba-ssm kernels (unimportable on ZeroGPU);
# False forces the in-tree NemotronH class with its native fallback.
self._tokenizer = AutoTokenizer.from_pretrained(self.model_id)
self._model = AutoModelForCausalLM.from_pretrained(
self.model_id,
dtype=torch.bfloat16, # v5 name (torch_dtype removed in transformers 5)
low_cpu_mem_usage=True,
)
self._model.eval()
# Move to cuda ONCE at load, guarded by `import spaces` (the lib
# virtualizes it and keeps the model resident across @spaces.GPU
# calls). With ONE-SHOT generation a petition is a single GPU call,
# so the multi-call NVML crash can't happen — and keeping the model
# resident avoids the ~67s CPU->GPU transfer that, done per-call,
# ate the whole GPU budget and left no time to generate (June 12).
# Resident + one call = ~14s real output.
try:
import spaces # noqa: F401
self._model.to("cuda")
except ImportError:
pass # local/CPU dev
def _generate(prompt: str, max_new_tokens: int, temperature: float) -> str:
import torch as _torch # noqa: F401
# Nemotron-Nano-9B-v2 is a reasoning model that THINKS by
# default; "/no_think" in the system slot disables it (June 12
# bench: raw completion leaked into outputs). The
# planner's full completion-style prompt rides as one user turn.
if getattr(self._tokenizer, "chat_template", None):
text_in = self._tokenizer.apply_chat_template(
[
{"role": "system", "content": "/no_think"},
{"role": "user", "content": prompt},
],
tokenize=False,
add_generation_prompt=True,
)
else:
text_in = prompt
inputs = self._tokenizer(
text_in,
return_tensors="pt",
truncation=True,
max_length=self.max_input_tokens,
).to(self._model.device)
with _torch.no_grad():
output = self._model.generate(
**inputs,
max_new_tokens=max_new_tokens,
do_sample=temperature > 0.05,
temperature=max(temperature, 0.05),
top_p=0.9,
pad_token_id=self._tokenizer.eos_token_id,
)
new_tokens = output[0][inputs["input_ids"].shape[1]:]
text = self._tokenizer.decode(new_tokens, skip_special_tokens=True)
# Defensive: strip any (possibly empty) think block that slips out.
text = re.sub(r".*?", "", text, flags=re.S)
text = text.replace("", "").replace("", "")
return text.strip()
try:
import spaces # import-guarded at call time, per contract
self._gpu_generate = spaces.GPU(duration=self.GPU_DURATION_S)(_generate)
except Exception:
self._gpu_generate = _generate
return self._gpu_generate
@staticmethod
def _is_valid(grammar: str, text: str) -> bool:
if is_moderation_grammar(grammar):
return parse_moderation(text)[1] is None
if base_tag(grammar) == FIX_TAG:
return parse_fix(text)[1] is None
return parse_build(text)[1] is None
async def _gen(self, generate, prompt: str, max_tokens: int, temperature: float) -> str:
"""Run one generation, retrying transient ZeroGPU CUDA/NVML failures.
Each @spaces.GPU call re-schedules onto a (possibly healthier) GPU, so a
retry can clear the intermittent NVML allocator assert seen June 12."""
loop = asyncio.get_running_loop()
last_exc = None
for attempt in range(3):
try:
return await loop.run_in_executor(None, generate, prompt, max_tokens, temperature)
except Exception as exc: # noqa: BLE001
msg = str(exc)
transient = any(s in msg for s in ("NVML", "CUDA", "cuda", "device", "out of memory"))
last_exc = exc
if not transient or attempt == 2:
raise
# drop the poisoned cuda placement so the next call re-homes the model
self._on_cuda = False
import torch as _t
try:
_t.cuda.empty_cache()
except Exception:
pass
await asyncio.sleep(0.5 * (attempt + 1))
raise last_exc # pragma: no cover
async def generate_stream(
self, prompt: str, grammar: str | None = None, max_tokens: int = 700
) -> AsyncIterator[str]:
loop = asyncio.get_running_loop()
generate = await loop.run_in_executor(None, self._ensure)
if grammar is None:
text = await self._gen(generate, prompt, max_tokens, 0.7)
else:
temperature = 0.0 if is_strict(grammar) else 0.4
strict_prompt = (
prompt + "\nReply with exactly one JSON object and nothing else."
)
text = await self._gen(generate, strict_prompt, max_tokens, temperature)
if not self._is_valid(grammar, text):
retry_prompt = (
strict_prompt
+ "\nYour previous output was malformed. Output ONLY the JSON object."
)
text = await self._gen(generate, retry_prompt, max_tokens, 0.0)
# Still malformed? Yield as-is; the planner retries then falls back.
for chunk in _chunks(text, 24):
await asyncio.sleep(0)
yield chunk
# --------------------------------------------------------------------------
# Factory
# --------------------------------------------------------------------------
_BACKEND_ALIASES = {
"mock": "mock",
"llamacpp": "llamacpp",
"llama": "llamacpp",
"gguf": "llamacpp",
"zerogpu": "zerogpu",
"zero-gpu": "zerogpu",
"transformers": "zerogpu",
}
def make_backend(name: str | None = None):
"""Build the backend selected by `name` or the CITY_BACKEND env.
Defaults to mock (always works, zero deps beyond stdlib).
"""
raw = (name or os.environ.get("CITY_BACKEND") or "mock").strip().lower()
resolved = _BACKEND_ALIASES.get(raw)
if resolved == "mock":
return MockBackend()
if resolved == "llamacpp":
return LlamaCppBackend()
if resolved == "zerogpu":
return ZeroGPUBackend()
raise ValueError(
f"unknown CITY_BACKEND {raw!r} (expected mock | zerogpu | llamacpp)"
)