Local LLM arena: invert a binary tree

Given the root of a binary tree, invert the tree and return its root.

Use this Python class for tree nodes:

    class TreeNode:
        def __init__(self, val=0, left=None, right=None):
            self.val = val
            self.left = left
            self.right = right

Implement a function with this exact signature:

    def invert_tree(root):
        ...

The argument is either a `TreeNode` or `None`. Return the root of the inverted tree.
You may mutate the input tree or build a new one — only the returned value is checked.

Output ONLY the function definition in a single Python code block.
Do not include explanations, example usage, tests, or print statements.

"""Send a prompt to a locally running llama-server, return the response.

Start the server first:
    llama-server -m <model.gguf> --port 8123 -c 8192 -ngl 99 -np 1
"""

import json
import re
import urllib.request


def call(prompt: str, *, port: int = 8123, temperature: float = 0, top_k: int = 1, seed: int = 42,
         enable_thinking: bool | None = None) -> tuple[str, dict]:
    """Returns (response_text, timings). `timings` is llama.cpp's per-call
    speed breakdown: predicted_n, predicted_per_second, predicted_ms, etc.

    `enable_thinking` forwards to the chat template via `chat_template_kwargs`
    so Qwen3-class models can be run with or without their default reasoning trace.
    Leave as None for models that don't have a thinking mode (e.g. Gemma)."""
    payload = {
        "messages": [{"role": "user", "content": prompt}],
        "temperature": temperature,
        "top_k": top_k,
        "seed": seed,
    }
    if enable_thinking is not None:
        payload["chat_template_kwargs"] = {"enable_thinking": enable_thinking}
    body = json.dumps(payload).encode()
    req = urllib.request.Request(
        f"http://127.0.0.1:{port}/v1/chat/completions",
        data=body,
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=300) as resp:
        data = json.loads(resp.read())
    return data["choices"][0]["message"]["content"], data.get("timings", {})


def extract_python(raw: str) -> str:
    m = re.search(r"```(?:python)?\s*\n(.*?)```", raw, re.DOTALL)
    return (m.group(1).rstrip() + "\n") if m else raw

"""Drive llama-server through each model in MODELS, write outputs to samples/.

llama-server's binary needs LD_LIBRARY_PATH pointing at its build/bin/ (libmtmd.so etc live next to it).
Model files are resolved via huggingface_hub, which downloads to ~/.cache/huggingface on first run.
"""

import json
import os
import subprocess
import sys
import threading
import time
import urllib.request
from pathlib import Path

from huggingface_hub import hf_hub_download

sys.path.insert(0, str(Path(__file__).parent))
import call_model


class Sampler:
    """Polls the target process (via /proc) and nvidia-smi while the context is open.
    Reports peak RAM (process RSS), CPU% across cores, peak VRAM, and GPU%."""

    def __init__(self, pid: int, interval: float = 0.2):
        self.pid = pid
        self.interval = interval
        self.gpu: list[tuple[float, float]] = []  # (util%, vram_mib)
        self.ram: list[float] = []                # process RSS in MiB
        self.cpu: list[float] = []                # process CPU% (can exceed 100 with multi-core)
        self._stop = threading.Event()
        self._thread: threading.Thread | None = None
        self._clk_tck = os.sysconf("SC_CLK_TCK")
        self._last_jiffies = 0
        self._last_time = 0.0

    def _sample_proc(self):
        try:
            rss_kib = 0
            with open(f"/proc/{self.pid}/status") as f:
                for line in f:
                    if line.startswith("VmRSS:"):
                        rss_kib = int(line.split()[1])
                        break
            with open(f"/proc/{self.pid}/stat") as f:
                fields = f.read().split()
            # utime + stime + cutime + cstime, in jiffies
            jiffies = sum(int(fields[i]) for i in (13, 14, 15, 16))
        except (FileNotFoundError, ProcessLookupError):
            return
        now = time.monotonic()
        self.ram.append(rss_kib / 1024)
        if self._last_jiffies and now > self._last_time:
            self.cpu.append((jiffies - self._last_jiffies) / self._clk_tck / (now - self._last_time) * 100)
        self._last_jiffies, self._last_time = jiffies, now

    def _sample_gpu(self):
        try:
            out = subprocess.check_output(
                ["nvidia-smi", "--query-gpu=utilization.gpu,memory.used",
                 "--format=csv,noheader,nounits"],
                text=True, timeout=1,
            )
            util, mem = (float(x) for x in out.strip().split(","))
            self.gpu.append((util, mem))
        except Exception:
            pass

    def _run(self):
        while not self._stop.is_set():
            self._sample_proc()
            self._sample_gpu()
            self._stop.wait(self.interval)

    def __enter__(self):
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()
        return self

    def __exit__(self, *exc):
        self._stop.set()
        if self._thread:
            self._thread.join(timeout=2)

    def summary(self) -> dict:
        avg = lambda xs: round(sum(xs) / len(xs), 1) if xs else None
        return {
            "ram_peak_mib":     int(max(self.ram)) if self.ram else None,
            "cpu_util_avg_pct": avg(self.cpu),
            "vram_peak_mib":    int(max(s[1] for s in self.gpu)) if self.gpu else None,
            "gpu_util_avg_pct": avg([s[0] for s in self.gpu]),
            "samples":          max(len(self.ram), len(self.gpu)),
        }

HERE = Path(__file__).parent
POST_DIR = HERE.parent
SAMPLES = POST_DIR / "verify" / "arena" / "samples"

LLAMA_BIN = Path.home() / ".unsloth" / "llama.cpp" / "build" / "bin" / "llama-server"

MODELS = [
    # (slug, repo, filename, enable_thinking, ngl).
    # enable_thinking=None: don't forward the flag (Gemma has no thinking mode).
    # enable_thinking=False on Qwen3 suppresses its default chain-of-thought trace,
    # which otherwise produces thousands of reasoning tokens before the answer.
    # ngl=99 offloads all layers to the GPU, ngl=0 runs purely on CPU.
    ("gemma-4-e4b-it-gpu", "unsloth/gemma-4-e4b-it-gguf", "gemma-4-E4B-it-UD-Q4_K_XL.gguf", None,  99),
    ("gemma-4-e4b-it-cpu", "unsloth/gemma-4-e4b-it-gguf", "gemma-4-E4B-it-UD-Q4_K_XL.gguf", None,  0),
    ("qwen3.5-4b-gpu",     "unsloth/Qwen3.5-4B-GGUF",     "Qwen3.5-4B-UD-Q4_K_XL.gguf",     False, 99),
    ("qwen3.5-4b-cpu",     "unsloth/Qwen3.5-4B-GGUF",     "Qwen3.5-4B-UD-Q4_K_XL.gguf",     False, 0),
]


def start_server(model_file, ngl: int):
    env = {**os.environ, "LD_LIBRARY_PATH": str(LLAMA_BIN.parent)}
    proc = subprocess.Popen(
        [str(LLAMA_BIN), "-m", str(model_file), "--port", "8123", "-c", "8192",
         "-ngl", str(ngl), "-np", "1", "--no-context-shift"],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, env=env,
    )
    for _ in range(240):
        try:
            urllib.request.urlopen("http://127.0.0.1:8123/health", timeout=1).read()
            return proc
        except Exception:
            time.sleep(0.5)
    proc.terminate()
    raise RuntimeError("llama-server did not become ready")


def stop_server(proc):
    proc.terminate()
    try:
        proc.wait(timeout=10)
    except subprocess.TimeoutExpired:
        proc.kill()


def main():
    prompt = (HERE / "prompt.md").read_text()
    for slug, repo, filename, enable_thinking, ngl in MODELS:
        print(f"--- {slug} ---", flush=True)
        model_file = hf_hub_download(repo_id=repo, filename=filename)
        proc = start_server(Path(model_file), ngl=ngl)
        try:
            with Sampler(pid=proc.pid) as s:
                raw, timings = call_model.call(prompt, enable_thinking=enable_thinking)
        finally:
            stop_server(proc)
        metrics = {
            "tokens_generated": timings.get("predicted_n"),
            "tokens_per_second": round(timings["predicted_per_second"], 1)
                if timings.get("predicted_per_second") else None,
            "generation_seconds": round(timings["predicted_ms"] / 1000, 2)
                if timings.get("predicted_ms") else None,
            **s.summary(),
        }
        out = SAMPLES / slug
        out.mkdir(parents=True, exist_ok=True)
        (out / "raw-response.txt").write_text(raw)
        (out / "solution.py").write_text(call_model.extract_python(raw))
        (out / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
        print(f"wrote {out.relative_to(POST_DIR)}/", flush=True)


if __name__ == "__main__":
    main()

def invert_tree(root):
    if root is None:
        return None
    
    # Swap the left and right children
    root.left, root.right = root.right, root.left
    
    # Recursively invert the subtrees
    invert_tree(root.left)
    invert_tree(root.right)
    
    return root

def invert_tree(root):
    if not root:
        return None
    
    root.left, root.right = root.right, root.left
    invert_tree(root.left)
    invert_tree(root.right)
    
    return root

"""Run each samples/<model>/solution.py against tree-inversion cases, print a pass-fail table."""

import importlib.util
import signal
from pathlib import Path


class T:
    def __init__(self, v, l=None, r=None):
        self.val, self.left, self.right = v, l, r


def equal(a, b):
    if a is None or b is None:
        return a is b
    return a.val == b.val and equal(a.left, b.left) and equal(a.right, b.right)


def cases():
    return [
        ("empty",       None,                                                 None),
        ("single",      T(1),                                                 T(1)),
        ("simple-3",    T(1, T(2), T(3)),                                     T(1, T(3), T(2))),
        ("balanced",    T(4, T(2, T(1), T(3)), T(7, T(6), T(9))),             T(4, T(7, T(9), T(6)), T(2, T(3), T(1)))),
        ("left-skewed", T(1, T(2, T(3))),                                     T(1, None, T(2, None, T(3)))),
    ]


def _alarm(*_):
    raise TimeoutError()

signal.signal(signal.SIGALRM, _alarm)


def evaluate(model_dir):
    spec = importlib.util.spec_from_file_location(model_dir.name, model_dir / "solution.py")
    mod = importlib.util.module_from_spec(spec)
    try:
        spec.loader.exec_module(mod)
        fn = mod.invert_tree
    except Exception:
        return ["-"] * 5
    marks = []
    for _, inp, expected in cases():
        signal.alarm(2)
        try:
            marks.append("✓" if equal(fn(inp), expected) else "✗")
        except Exception:
            marks.append("✗")
        finally:
            signal.alarm(0)
    return marks


def main():
    samples = Path(__file__).parent / "samples"
    rows = sorted((d.name, evaluate(d)) for d in samples.iterdir() if d.is_dir())
    names = [c[0] for c in cases()]
    print("| model | " + " | ".join(names) + " | total |")
    print("|" + "---|" * (len(names) + 2))
    for name, marks in rows:
        print(f"| {name} | " + " | ".join(marks) + f" | {sum(m == '✓' for m in marks)}/{len(names)} |")


if __name__ == "__main__":
    main()