Inference on NRP

import os, shutil, subprocess

print("OPENAI_API_BASE =", os.environ.get("OPENAI_API_BASE", "NOT SET"))
# Print the actual token. This is a shared workshop key — fine to show here;
# treat your own personal token as a secret.
key = os.environ.get("OPENAI_API_KEY", "")
print("OPENAI_API_KEY  =", key if key else "NOT SET")

kubectl = shutil.which("kubectl") or "/opt/conda/bin/kubectl"
print("kubectl path    =", kubectl)
print("kubectl version =", subprocess.run([kubectl, "version", "--client", "-o", "json"],
                                          capture_output=True, text=True).stdout[:120], "...")

print("can list pods in nrp-training-k8s:",
      subprocess.run([kubectl, "auth", "can-i", "list", "pods", "-n", "nrp-training-k8s"],
                     capture_output=True, text=True).stdout.strip())

# GPU detection. Section 3 (local Ollama) need this; others are fine without.
import subprocess

HAS_GPU = False
try:
    out = subprocess.run(["nvidia-smi", "-L"], capture_output=True, text=True, check=True).stdout.strip()
    print("GPU(s) detected:")
    print(out)
    HAS_GPU = True
except (FileNotFoundError, subprocess.CalledProcessError):
    print("No GPU in this session pod.")
    print("Section 3 (local Ollama) will skip the local-LLM half.")
    print("To run the full comparison, respawn with 1 x NVIDIA-A10 from the JupyterHub launcher.")

# List the models currently served by the NRP managed endpoint.
import os, json
from openai import OpenAI

client = OpenAI(api_key=os.environ["OPENAI_API_KEY"], base_url=os.environ["OPENAI_API_BASE"])
models = client.models.list()
print(f"{len(models.data)} models available right now:\n")
for m in sorted(models.data, key=lambda x: x.id):
    print(f"  {m.id}")

# A tiny helper we'll reuse across this notebook. The default model is
# `gemma-small-e4b` (Gemma 3n E4B — fast on NRP and multimodal). The helper also
# handles *reasoning* models like `minimax-m2` (you'll try one below): they
# stream a private chain-of-thought into a separate `reasoning` field and only
# fill `content` once they conclude — so we default `max_tokens` high enough
# (1200) to leave room for both the thinking and the final answer. If a call
# still runs out of tokens mid-thought, we surface that reasoning rather than
# returning nothing.
def chat(prompt, model="gemma-small-e4b", system=None, max_tokens=1200, llm=None):
    msgs = []
    if system:
        msgs.append({"role": "system", "content": system})
    msgs.append({"role": "user", "content": prompt})
    resp = (llm or client).chat.completions.create(
        model=model, messages=msgs, max_tokens=max_tokens, temperature=0.2,
    )
    msg = resp.choices[0].message
    if msg.content:
        return msg.content
    reasoning = getattr(msg, "reasoning", None) or getattr(msg, "reasoning_content", None)
    if reasoning:
        return f"(no final answer within max_tokens={max_tokens}; reasoning so far:)\n{reasoning}"
    return "(model produced no content — try raising max_tokens or another model)"

# Single chat completion.
print(chat(
    "What is the National Research Platform?",
    system="Answer in one sentence.",
))

# Streaming — tokens arrive live, one at a time, instead of all at once.
# minimax-m2 is a reasoning model: it streams its private "thinking" first (the
# `reasoning` field), then the final answer (`content`). We print BOTH, labeled,
# so you can watch the tokens flow and see what "reasoning" means. (A loop that
# printed only `content` — like a haiku-only loop — would show nothing until the
# thinking finished, and with a small max_tokens it may never finish.)
stream = client.chat.completions.create(
    model="minimax-m2",
    messages=[{"role": "user", "content": "Write a haiku about GPUs."}],
    stream=True,
    max_tokens=600,
)

section = None
for chunk in stream:
    if not chunk.choices:
        continue
    delta = chunk.choices[0].delta
    think = getattr(delta, "reasoning", None) or getattr(delta, "reasoning_content", None)
    if think:
        if section != "think":
            print("[thinking]"); section = "think"
        print(think, end="", flush=True)
    if delta.content:
        if section != "answer":
            print("\n\n[answer]"); section = "answer"
        print(delta.content, end="", flush=True)
print()

# Switch to a smaller model — same code, just change the `model` arg.
print(chat("Explain Kubernetes namespaces in two sentences.", model="gemma-small-e4b"))

# Now a reasoning model. We give it room (max_tokens=800) so it finishes
# thinking AND produces a final answer — with too few tokens the whole budget
# is spent reasoning and `content` comes back empty (what the helper guards).
print("\n--- minimax-m2 (a reasoning model) ---")
print(chat("In one sentence, name a strength of this model.", model="minimax-m2", max_tokens=800))

# Try the others yourself — e.g. model="gpt-oss" (fast) or model="qwen3"
# (the 397B flagship; highest quality but slowest to respond).

# Install Ollama onto the session's FAST local disk (/tmp), NOT the home
# directory. Ollama ships ~3.5 GB of CUDA libraries; the home volume is
# networked storage (Ceph), and reading 3.5 GB from it is too slow for Ollama's
# GPU probe — it times out and silently falls back to CPU. /tmp is node-local
# and fast, so the A10 is detected in under a second. (/tmp is wiped when the
# session stops, so this re-downloads each session — it only takes ~10 s.)
import os, shutil, subprocess

OLLAMA_DIR = f"/tmp/ollama-{os.getuid()}"  # per-user dir we own (avoids shared-path clashes)
OLLAMA_BIN = f"{OLLAMA_DIR}/bin/ollama"
OLLAMA_VERSION = "v0.24.0"  # pinned to a known-good release

if not HAS_GPU:
    print("Skipping Ollama install — no GPU in this session.")
elif os.path.exists(OLLAMA_BIN):
    print("Ollama already installed at", OLLAMA_BIN)
else:
    print(f"Downloading Ollama {OLLAMA_VERSION} to {OLLAMA_DIR} (~10 s)...")
    os.makedirs(OLLAMA_DIR, exist_ok=True)
    url = f"https://github.com/ollama/ollama/releases/download/{OLLAMA_VERSION}/ollama-linux-amd64.tar.zst"
    subprocess.run(
        f"curl -fsSL {url} | tar --use-compress-program=unzstd -xf - -C {OLLAMA_DIR}",
        shell=True, check=True,
    )
if HAS_GPU:
    os.environ["PATH"] = f"{OLLAMA_DIR}/bin" + os.pathsep + os.environ["PATH"]
    print("ollama binary at:", shutil.which("ollama"))

# Start `ollama serve` in the background and wait until it answers.
#   - Everything (binary, CUDA libs, models) lives under /tmp/ol on fast local
#     disk, so the GPU probe finds the A10 in well under a second.
#   - OLLAMA_LLM_LIBRARY + LD_LIBRARY_PATH point at the bundled CUDA libraries.
#   - OLLAMA_MODELS keeps pulled models on /tmp too, so they load fast.
import os, socket, subprocess, time, urllib.request, urllib.error

def ollama_alive():
    try:
        urllib.request.urlopen("http://localhost:11434/api/tags", timeout=3)
        return True
    except (urllib.error.URLError, ConnectionResetError, OSError, socket.timeout, TimeoutError):
        return False

LIB = f"{OLLAMA_DIR}/lib/ollama"
SERVE_LOG = f"{OLLAMA_DIR}/serve.log"  # we created OLLAMA_DIR, so this is always writable

if not HAS_GPU:
    print("Skipping Ollama serve — no GPU in this session.")
elif ollama_alive():
    print("Ollama already running.")
else:
    print("Starting `ollama serve` in the background...")
    ollama_env = {
        **os.environ,
        "HOME": OLLAMA_DIR,  # keep ollama's identity key + data on the fast local dir we own
        "OLLAMA_HOST": "0.0.0.0:11434",
        "OLLAMA_LOAD_TIMEOUT": "15m",
        "OLLAMA_MODELS": f"{OLLAMA_DIR}/models",
        "OLLAMA_LIBRARY_PATH": LIB,
        "OLLAMA_LLM_LIBRARY": "cuda_v12",
        "LD_LIBRARY_PATH": ":".join([
            LIB, LIB + "/cuda_v12", "/usr/lib/x86_64-linux-gnu",
            os.environ.get("LD_LIBRARY_PATH", ""),
        ]),
    }
    subprocess.Popen(f"nohup {OLLAMA_BIN} serve > {SERVE_LOG} 2>&1 &", shell=True, env=ollama_env)
    t0 = time.time(); deadline = t0 + 120
    while time.time() < deadline:
        time.sleep(2)
        if ollama_alive():
            print(f"Ollama up after ~{int(time.time() - t0)}s.")
            break
    else:
        print(f"Ollama did not come up within 120s. Check {SERVE_LOG}")

# Report which device Ollama is using.
if HAS_GPU and ollama_alive():
    log = subprocess.run(f"grep -iE 'inference compute' {SERVE_LOG} | tail -1",
                         shell=True, capture_output=True, text=True).stdout.lower()
    if "library=cuda" in log:
        print("Device: GPU (CUDA) — the A10 was detected. \U0001f389")
    elif "id=cpu" in log:
        print(f"Device: CPU — GPU probe fell back to CPU; check {SERVE_LOG}")
    else:
        print(f"Device: unknown — see {SERVE_LOG}")

# Pull the mistral 7B model (~4 GB). First time only.
import subprocess, time, json, urllib.request

def already_pulled(name="mistral"):
    try:
        r = json.load(urllib.request.urlopen("http://localhost:11434/api/tags", timeout=2))
        return any(name in m["name"] for m in r.get("models", []))
    except Exception:
        return False

if not HAS_GPU:
    print("Skipping — no GPU.")
elif already_pulled("mistral"):
    print("mistral already pulled.")
else:
    print("Pulling mistral... (~4 GB, 2-4 minutes the first time)")
    # Drive the pull via the HTTP API instead of `ollama pull`, because the CLI
    # writes \r-based progress bars that buffer indefinitely under nbconvert.
    body = json.dumps({"name": "mistral", "stream": True}).encode()
    req  = urllib.request.Request("http://localhost:11434/api/pull", data=body,
                                  headers={"Content-Type": "application/json"})
    t0 = time.time()
    last_status = None
    with urllib.request.urlopen(req) as resp:
        for raw in resp:
            try:
                evt = json.loads(raw.decode())
            except Exception:
                continue
            status = evt.get("status", "")
            # Only print on status transitions (not every kilobyte) to keep output tidy
            if status and status != last_status:
                print(f"  [{int(time.time()-t0):3d}s] {status}")
                last_status = status
            if evt.get("error"):
                raise RuntimeError(evt["error"])
    print(f"Done in {int(time.time()-t0)}s.")

# Hello-world against the local model. Same OpenAI client, different base_url.
# Heads up: the first call after `ollama serve` started will be slow (~30-120s
# while llama.cpp loads the 4 GB of weights into VRAM). Subsequent calls in
# this notebook reuse the loaded model and respond in 1-2 seconds.
from openai import OpenAI

if HAS_GPU:
    local = OpenAI(api_key="ollama", base_url="http://localhost:11434/v1",
                   timeout=900)  # generous timeout for the cold load
    print(chat(
        "Say hi from the GPU in this JupyterHub session.",
        model="mistral", llm=local, max_tokens=80,
    ))
else:
    print("Skipping — no GPU.")

# Side-by-side: same prompt, managed minimax-m2 vs local mistral.
prompt = "Explain in two sentences why a research cluster might use Kubernetes namespaces."

print("=" * 78)
print("NRP MANAGED (minimax-m2)")
print("=" * 78)
print(chat(prompt, model="minimax-m2"))

if HAS_GPU:
    print()
    print("=" * 78)
    print("LOCAL (mistral, on this pod's GPU)")
    print("=" * 78)
    print(chat(prompt, model="mistral", llm=local))
else:
    print("\n(Local half skipped — no GPU.)")

# Pull one real NRP documentation page (the cluster policies page) and clean it.
import requests, re

PAGE = "https://nrp.ai/documentation/userdocs/start/policies/"
RAW  = ("https://gitlab.nrp-nautilus.io/prp/nrp-site/-/raw/main/"
        "src/content/docs/Documentation/userdocs/start/policies.mdx")

md = requests.get(RAW, timeout=30).text
md = re.sub(r"^---.*?---\s*", "", md, flags=re.S)        # drop frontmatter
md = re.sub(r"^import .*$", "", md, flags=re.M)          # drop MDX imports
md = re.sub(r":::\w+(\[[^\]]*\])?|:::", "", md)          # drop admonition markers
md = re.sub(r"<[^>]+>", "", md)                          # drop stray HTML/JSX
md = re.sub(r"\n{3,}", "\n\n", md).strip()

# Split into ~700-char chunks. (A real corpus would chunk every page; one is plenty here.)
chunks = [md[i:i+700] for i in range(0, len(md), 550)]
print(f"Pulled {len(md)} chars from {PAGE}")
print(f"Split into {len(chunks)} chunks.")

# Embed with NRP's MANAGED embedding model — no local model, nothing to download.
# Same `client`, same endpoint as the chat models; just a different call.
import numpy as np

def embed(texts):
    resp = client.embeddings.create(model="qwen3-embedding", input=texts)
    vecs = np.array([d.embedding for d in resp.data])
    return vecs / np.linalg.norm(vecs, axis=1, keepdims=True)   # normalize for cosine

chunk_vecs = embed(chunks)          # one API call embeds the whole page
print(f"Embedded {len(chunks)} chunks into vectors of dim {chunk_vecs.shape[1]} (via qwen3-embedding).")

def retrieve(question, k=3):
    qv = embed([question])[0]
    scores = chunk_vecs @ qv         # cosine similarity (everything is normalized)
    top = scores.argsort()[-k:][::-1]
    return [(chunks[i], float(scores[i])) for i in top]

# Answer a question by RAG. Swap model/llm to answer with the local GPU model instead.
SYSTEM = ("Answer the question using ONLY the provided context. "
          "If the context doesn't contain the answer, say so. Be concise.")

def ask(question, model="minimax-m2", llm=None):
    context = "\n\n".join(text for text, _ in retrieve(question))
    return chat(f"Context:\n{context}\n\nQuestion: {question}", system=SYSTEM, model=model, llm=llm)

QUESTION = "Is it okay to run `sleep infinity` in a Job on NRP?"
print("Retrieved chunks:")
for text, score in retrieve(QUESTION):
    print(f"  score={score:.3f}  {text[:70].strip()}...")
print("\n--- Answer (NRP managed minimax-m2) ---")
print(ask(QUESTION))

# The exact same RAG pipeline, answered by the local model on this session's GPU.
# Only the model + client change — retrieval and embeddings are identical.
if HAS_GPU:
    print("--- Answer (local mistral on the GPU) ---")
    print(ask(QUESTION, model="mistral", llm=local))
else:
    print("(Skipped — no GPU. The managed answer above used the same pipeline.)")

# Stop the local Ollama daemon (frees GPU memory).
import subprocess
subprocess.run(["pkill", "-x", "ollama"], check=False)
print("Ollama stopped.")