TinyCodeGPT：从零训练一个本地任务代码生成模型TinyCodeGPT: Train a Local Task-to-Code LLM from Scratch

from __future__ import annotations

import argparse
import base64
import json
import math
import os
import random
import subprocess
import sys
import tempfile
import threading
import time
import traceback
import uuid
from dataclasses import dataclass
from datetime import datetime, timezone
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any
from urllib.parse import urlparse

try:
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
except Exception:  # pragma: no cover - status endpoint reports this clearly.
    torch = None
    class _NNStub:
        Module = object

    nn = _NNStub()
    F = None


def torch_no_grad():
    if torch is not None:
        return torch.no_grad()

    def decorator(function):
        return function

    return decorator


ROOT_DIR = Path(__file__).resolve().parents[1]
RUN_DIR = ROOT_DIR / ".tmp" / "tinycodegpt"
CHECKPOINT_DIR = RUN_DIR / "checkpoints"
RUN_OUTPUT_DIR = RUN_DIR / "generated_runs"
DATASET_DIR = RUN_DIR / "datasets"
RUN_DIR.mkdir(parents=True, exist_ok=True)
CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
RUN_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
DATASET_DIR.mkdir(parents=True, exist_ok=True)

JOBS: dict[str, dict[str, Any]] = {}
STOP_EVENTS: dict[str, threading.Event] = {}
ACTIVE_JOB_ID: str | None = None
LOCK = threading.Lock()


CATEGORIES = ("Math", "Python", "NumPy", "Pandas", "Matplotlib", "ML")
FORBIDDEN_CODE_PATTERNS = (
    "subprocess",
    "os.system",
    "socket",
    "requests",
    "urllib",
    "shutil",
    "pickle",
    "eval(",
    "exec(",
    "open(",
    "Path(",
)


@dataclass
class ModelConfig:
    n_layer: int
    n_head: int
    n_embd: int
    block_size: int


PRESETS = {
    "tiny": ModelConfig(n_layer=2, n_head=2, n_embd=96, block_size=192),
    "small": ModelConfig(n_layer=4, n_head=4, n_embd=192, block_size=320),
    "medium": ModelConfig(n_layer=6, n_head=6, n_embd=384, block_size=512),
}


def clamp_int(value: Any, fallback: int, lower: int, upper: int) -> int:
    try:
        parsed = int(value)
    except (TypeError, ValueError):
        parsed = fallback
    return max(lower, min(parsed, upper))


def clamp_float(value: Any, fallback: float, lower: float, upper: float) -> float:
    try:
        parsed = float(value)
    except (TypeError, ValueError):
        parsed = fallback
    return max(lower, min(parsed, upper))


def resolve_model_config(request: dict[str, Any]) -> tuple[str, ModelConfig]:
    preset = str(request.get("preset") or "small")
    base_config = PRESETS.get(preset, PRESETS["small"])
    raw_config = request.get("model_config") if preset == "custom" else None
    if not isinstance(raw_config, dict):
        return preset if preset in PRESETS else "small", base_config

    config = ModelConfig(
        n_layer=clamp_int(raw_config.get("n_layer"), base_config.n_layer, 1, 12),
        n_head=clamp_int(raw_config.get("n_head"), base_config.n_head, 1, 12),
        n_embd=clamp_int(raw_config.get("n_embd"), base_config.n_embd, 32, 768),
        block_size=clamp_int(raw_config.get("block_size"), base_config.block_size, 64, 1024),
    )
    if config.n_embd % config.n_head != 0:
        raise ValueError("n_embd must be divisible by n_head for custom model config.")
    return "custom", config


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


def jsonable(value: Any) -> Any:
    if isinstance(value, Path):
        return str(value)
    return value


def respond(handler: BaseHTTPRequestHandler, status: int, payload: dict[str, Any]) -> None:
    raw = json.dumps(payload, ensure_ascii=False, default=jsonable).encode("utf-8")
    handler.send_response(status)
    handler.send_header("Content-Type", "application/json; charset=utf-8")
    handler.send_header("Content-Length", str(len(raw)))
    handler.send_header("Access-Control-Allow-Origin", "*")
    handler.send_header("Access-Control-Allow-Headers", "Content-Type")
    handler.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
    handler.end_headers()
    handler.wfile.write(raw)


def read_body(handler: BaseHTTPRequestHandler) -> dict[str, Any]:
    length = int(handler.headers.get("Content-Length") or "0")
    if length <= 0:
        return {}
    raw = handler.rfile.read(length).decode("utf-8")
    return json.loads(raw or "{}")


def update_job(job_id: str, **patch: Any) -> None:
    with LOCK:
        job = JOBS[job_id]
        job.update(patch)


def append_log(job_id: str, line: str) -> None:
    with LOCK:
        job = JOBS[job_id]
        logs = job.setdefault("logs", [])
        logs.append(line)
        if len(logs) > 240:
            del logs[: len(logs) - 240]


def append_metric(job_id: str, metric: dict[str, Any]) -> None:
    with LOCK:
        job = JOBS[job_id]
        metrics = job.setdefault("metrics", [])
        metrics.append(metric)
        if len(metrics) > 240:
            del metrics[: len(metrics) - 240]


def add_sample(samples: list[dict[str, str]], category: str, task: str, code: str) -> None:
    samples.append({"category": category, "task": task.strip(), "code": code.strip()})


def generate_dataset(samples_per_category: int, seed: int = 7) -> list[dict[str, str]]:
    rng = random.Random(seed)
    samples: list[dict[str, str]] = []

    for _ in range(samples_per_category):
        r = rng.randint(1, 80)
        add_sample(
            samples,
            "Math",
            f"calculate the area of a circle with radius {r}",
            f"import math\nr = {r}\narea = math.pi * r ** 2\nprint(area)",
        )
        w, h = rng.randint(2, 60), rng.randint(2, 60)
        add_sample(
            samples,
            "Math",
            f"calculate rectangle area with width {w} and height {h}",
            f"width = {w}\nheight = {h}\narea = width * height\nprint(area)",
        )
        numbers = [rng.randint(-20, 80) for _ in range(rng.randint(4, 8))]
        add_sample(
            samples,
            "Math",
            f"calculate mean and variance of {numbers}",
            "nums = "
            + repr(numbers)
            + "\nmean = sum(nums) / len(nums)\nvariance = sum((x - mean) ** 2 for x in nums) / len(nums)\nprint(mean, variance)",
        )
        x1, y1, x2, y2 = [rng.randint(-20, 20) for _ in range(4)]
        add_sample(
            samples,
            "Math",
            f"calculate distance between ({x1}, {y1}) and ({x2}, {y2})",
            f"import math\nx1, y1 = {x1}, {y1}\nx2, y2 = {x2}, {y2}\ndistance = math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)\nprint(distance)",
        )

        py_nums = [rng.randint(-30, 90) for _ in range(rng.randint(5, 10))]
        add_sample(
            samples,
            "Python",
            f"sort the list {py_nums} in ascending order",
            f"nums = {py_nums}\nsorted_nums = sorted(nums)\nprint(sorted_nums)",
        )
        threshold = rng.randint(0, 40)
        add_sample(
            samples,
            "Python",
            f"filter numbers greater than {threshold} from {py_nums}",
            f"nums = {py_nums}\nfiltered = [x for x in nums if x > {threshold}]\nprint(filtered)",
        )
        words = rng.choices(["alpha", "beta", "gamma", "delta", "theta", "lambda"], k=8)
        add_sample(
            samples,
            "Python",
            f"count word frequencies in {' '.join(words)}",
            "text = "
            + repr(" ".join(words))
            + "\ncounts = {}\nfor word in text.split():\n    counts[word] = counts.get(word, 0) + 1\nprint(counts)",
        )

        arr = [rng.randint(1, 50) for _ in range(5)]
        add_sample(
            samples,
            "NumPy",
            f"normalize numpy array {arr} to range 0 to 1",
            "import numpy as np\narr = np.array("
            + repr(arr)
            + ", dtype=float)\nnormalized = (arr - arr.min()) / (arr.max() - arr.min())\nprint(normalized)",
        )
        a = [[rng.randint(1, 5), rng.randint(1, 5)], [rng.randint(1, 5), rng.randint(1, 5)]]
        b = [[rng.randint(1, 5), rng.randint(1, 5)], [rng.randint(1, 5), rng.randint(1, 5)]]
        add_sample(
            samples,
            "NumPy",
            f"multiply two numpy matrices {a} and {b}",
            f"import numpy as np\nA = np.array({a})\nB = np.array({b})\nC = A @ B\nprint(C)",
        )
        start, stop = rng.randint(-10, 0), rng.randint(5, 20)
        add_sample(
            samples,
            "NumPy",
            f"create {8} evenly spaced numpy values from {start} to {stop}",
            f"import numpy as np\nx = np.linspace({start}, {stop}, 8)\nprint(x)",
        )

        prices = [rng.randint(10, 100) for _ in range(6)]
        cities = rng.choices(["Beijing", "Shanghai", "Chengdu"], k=6)
        add_sample(
            samples,
            "Pandas",
            "calculate mean price by city in a pandas dataframe",
            "import pandas as pd\n"
            + f"df = pd.DataFrame({{'city': {cities}, 'price': {prices}}})\n"
            + "result = df.groupby('city')['price'].mean()\nprint(result)",
        )
        add_sample(
            samples,
            "Pandas",
            f"filter pandas rows where price is greater than {threshold}",
            "import pandas as pd\n"
            + f"df = pd.DataFrame({{'city': {cities}, 'price': {prices}}})\n"
            + f"filtered = df[df['price'] > {threshold}]\nprint(filtered)",
        )

        func = rng.choice(["x ** 2", "np.sin(x)", "np.cos(x)", "x ** 3"])
        add_sample(
            samples,
            "Matplotlib",
            f"plot y = {func} from -10 to 10",
            "import numpy as np\nimport matplotlib.pyplot as plt\nx = np.linspace(-10, 10, 100)\ny = "
            + func
            + "\nplt.plot(x, y)\nplt.title('function plot')\nplt.show()",
        )
        add_sample(
            samples,
            "Matplotlib",
            "draw a scatter plot for two small numeric lists",
            "import matplotlib.pyplot as plt\nx = [1, 2, 3, 4, 5]\ny = [2, 5, 4, 8, 7]\nplt.scatter(x, y)\nplt.xlabel('x')\nplt.ylabel('y')\nplt.show()",
        )

        slope = rng.uniform(0.5, 4.0)
        intercept = rng.uniform(-3.0, 3.0)
        xs = list(range(-5, 6))
        ys = [round(slope * x + intercept, 3) for x in xs]
        add_sample(
            samples,
            "ML",
            "fit linear regression with numpy least squares",
            "import numpy as np\n"
            + f"x = np.array({xs}, dtype=float)\ny = np.array({ys}, dtype=float)\n"
            + "X = np.c_[x, np.ones_like(x)]\ncoef, bias = np.linalg.lstsq(X, y, rcond=None)[0]\nprint(coef, bias)",
        )
        add_sample(
            samples,
            "ML",
            "run gradient descent for y = 2x + 1",
            "import numpy as np\nx = np.array([0, 1, 2, 3], dtype=float)\ny = 2 * x + 1\nw, b = 0.0, 0.0\nlr = 0.05\nfor step in range(200):\n    pred = w * x + b\n    loss = np.mean((pred - y) ** 2)\n    dw = np.mean(2 * (pred - y) * x)\n    db = np.mean(2 * (pred - y))\n    w -= lr * dw\n    b -= lr * db\nprint(round(w, 3), round(b, 3), round(loss, 6))",
        )

    rng.shuffle(samples)
    return samples


def strip_jsonl_fences(raw: str) -> str:
    text = raw.strip()
    if text.startswith("```"):
        lines = text.splitlines()
        if lines and lines[0].lstrip().startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        text = "\n".join(lines).strip()
    return text


def assert_safe_sample_code(code: str, line_number: int) -> None:
    lowered = code.lower()
    for pattern in FORBIDDEN_CODE_PATTERNS:
        if pattern.lower() in lowered:
            raise ValueError(f"line {line_number}: code contains forbidden operation {pattern!r}")


def normalize_dataset_sample(value: Any, line_number: int) -> dict[str, str]:
    if not isinstance(value, dict):
        raise ValueError(f"line {line_number}: expected a JSON object")
    category = str(value.get("category") or "").strip()
    task = str(value.get("task") or "").strip()
    code = str(value.get("code") or "").strip()
    if category not in CATEGORIES:
        raise ValueError(f"line {line_number}: category must be one of {', '.join(CATEGORIES)}")
    if not task:
        raise ValueError(f"line {line_number}: task is required")
    if not code:
        raise ValueError(f"line {line_number}: code is required")
    assert_safe_sample_code(code, line_number)
    return {"category": category, "task": task, "code": code}


def parse_jsonl_dataset(raw: str) -> list[dict[str, str]]:
    text = strip_jsonl_fences(raw)
    if not text:
        raise ValueError("dataset_jsonl is empty")
    if text.startswith("["):
        parsed = json.loads(text)
        if not isinstance(parsed, list):
            raise ValueError("dataset JSON array is invalid")
        samples = [normalize_dataset_sample(item, index + 1) for index, item in enumerate(parsed)]
    else:
        samples = []
        for index, line in enumerate(text.splitlines(), start=1):
            stripped = line.strip()
            if not stripped:
                continue
            samples.append(normalize_dataset_sample(json.loads(stripped), index))
    if not samples:
        raise ValueError("dataset_jsonl does not contain any samples")
    return samples


def save_jsonl_dataset(job_id: str, samples: list[dict[str, str]]) -> Path:
    dataset_path = DATASET_DIR / f"{job_id}.jsonl"
    payload = "\n".join(json.dumps(sample, ensure_ascii=False) for sample in samples) + "\n"
    dataset_path.write_text(payload, encoding="utf-8")
    return dataset_path


def load_training_samples(
    job_id: str,
    request: dict[str, Any],
    samples_per_category: int,
    seed: int,
) -> tuple[list[dict[str, str]], str, Path | None]:
    raw_jsonl = str(request.get("dataset_jsonl") or "").strip()
    if raw_jsonl:
        samples = parse_jsonl_dataset(raw_jsonl)
        dataset_path = save_jsonl_dataset(job_id, samples)
        return samples, "teacher_jsonl", dataset_path

    return generate_dataset(samples_per_category, seed=seed), "template_synthetic", None


def format_sample(sample: dict[str, str]) -> str:
    return f"<task>\n{sample['task']}\n</task>\n<code>\n{sample['code']}\n</code>\n"


def build_vocab(text: str) -> tuple[dict[str, int], dict[int, str]]:
    chars = sorted(set(text))
    stoi = {ch: i for i, ch in enumerate(chars)}
    itos = {i: ch for ch, i in stoi.items()}
    return stoi, itos


def count_parameters(model: nn.Module) -> int:
    return sum(parameter.numel() for parameter in model.parameters())


def normalize_itos(raw_itos: dict[Any, str]) -> dict[int, str]:
    if not raw_itos:
        return {}
    return {int(key): value for key, value in raw_itos.items()} if isinstance(next(iter(raw_itos.keys())), str) else raw_itos


def load_resume_checkpoint(path_value: str | None) -> tuple[Path, dict[str, Any]]:
    checkpoint_path = Path(path_value) if path_value else latest_checkpoint()
    if checkpoint_path is None or not checkpoint_path.exists():
        raise RuntimeError("No checkpoint found to continue training from.")
    return checkpoint_path, torch.load(checkpoint_path, map_location="cpu")


class CausalSelfAttention(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        if config.n_embd % config.n_head != 0:
            raise ValueError("n_embd must be divisible by n_head")
        self.n_head = config.n_head
        self.head_dim = config.n_embd // config.n_head
        self.qkv = nn.Linear(config.n_embd, 3 * config.n_embd)
        self.proj = nn.Linear(config.n_embd, config.n_embd)
        self.register_buffer("mask", torch.tril(torch.ones(config.block_size, config.block_size)).view(1, 1, config.block_size, config.block_size))

    def forward(self, x):
        batch, tokens, channels = x.shape
        qkv = self.qkv(x)
        q, k, v = qkv.split(channels, dim=2)
        q = q.view(batch, tokens, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(batch, tokens, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(batch, tokens, self.n_head, self.head_dim).transpose(1, 2)
        attention = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        attention = attention.masked_fill(self.mask[:, :, :tokens, :tokens] == 0, float("-inf"))
        attention = F.softmax(attention, dim=-1)
        y = attention @ v
        y = y.transpose(1, 2).contiguous().view(batch, tokens, channels)
        return self.proj(y)


class TransformerBlock(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.ln1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln2 = nn.LayerNorm(config.n_embd)
        self.ffn = nn.Sequential(
            nn.Linear(config.n_embd, 4 * config.n_embd),
            nn.GELU(),
            nn.Linear(4 * config.n_embd, config.n_embd),
        )

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x


class TinyCodeGPT(nn.Module):
    def __init__(self, vocab_size: int, config: ModelConfig):
        super().__init__()
        self.config = config
        self.token_emb = nn.Embedding(vocab_size, config.n_embd)
        self.pos_emb = nn.Embedding(config.block_size, config.n_embd)
        self.blocks = nn.Sequential(*[TransformerBlock(config) for _ in range(config.n_layer)])
        self.ln_f = nn.LayerNorm(config.n_embd)
        self.head = nn.Linear(config.n_embd, vocab_size, bias=False)

    def forward(self, idx, targets=None):
        _, tokens = idx.shape
        positions = torch.arange(tokens, device=idx.device)
        x = self.token_emb(idx) + self.pos_emb(positions)
        x = self.blocks(x)
        logits = self.head(self.ln_f(x))
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
        return logits, loss

    @torch_no_grad()
    def generate(
        self,
        idx,
        max_new_tokens: int,
        temperature: float = 0.8,
        top_k: int = 30,
        stop_sequence: list[int] | None = None,
    ):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.config.block_size :]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / max(temperature, 1e-4)
            if top_k:
                values, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < values[:, [-1]]] = -float("inf")
            probs = F.softmax(logits, dim=-1)
            next_id = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, next_id), dim=1)
            if stop_sequence and idx.size(1) >= len(stop_sequence):
                tail = idx[0, -len(stop_sequence) :].detach().cpu().tolist()
                if tail == stop_sequence:
                    break
        return idx


def make_batch(data: Any, block_size: int, batch_size: int, device: str):
    starts = torch.randint(0, len(data) - block_size - 1, (batch_size,))
    x = torch.stack([data[i : i + block_size] for i in starts]).to(device)
    y = torch.stack([data[i + 1 : i + block_size + 1] for i in starts]).to(device)
    return x, y


def evaluate(model: nn.Module, data: Any, block_size: int, batch_size: int, device: str) -> float:
    model.eval()
    losses = []
    with torch.no_grad():
        for _ in range(6):
            x, y = make_batch(data, block_size, batch_size, device)
            _, loss = model(x, y)
            losses.append(float(loss.item()))
    model.train()
    return sum(losses) / len(losses)


def run_training(job_id: str, request: dict[str, Any]) -> None:
    global ACTIVE_JOB_ID
    if torch is None:
        update_job(job_id, status="failed", error="PyTorch is not installed in this Python environment.", finished_at=utc_now())
        return

    stop_event = STOP_EVENTS[job_id]
    try:
        continue_from_checkpoint = bool(request.get("continue_from_checkpoint"))
        resume_checkpoint_path = str(request.get("checkpoint_path") or "").strip() or None
        preset, config = resolve_model_config(request)
        samples_per_category = max(10, min(int(request.get("samples_per_category") or 80), 600))
        fallback_max_steps = max(20, min(int(request.get("max_steps") or 50_000), 200_000))
        target_epochs = clamp_float(request.get("target_epochs"), 3.0, 0.0, 100.0)
        batch_size = max(4, min(int(request.get("batch_size") or 32), 128))
        learning_rate = float(request.get("learning_rate") or 8e-4)
        seed = int(request.get("seed") or 7)

        random.seed(seed)
        torch.manual_seed(seed)
        device = "cuda" if torch.cuda.is_available() else "cpu"
        update_job(job_id, status="running", started_at=utc_now(), device=device)
        append_log(job_id, f"device = {device}")
        append_log(job_id, f"preset = {preset}, config = {config}")

        samples, dataset_source, dataset_path = load_training_samples(job_id, request, samples_per_category, seed)
        text = "\n".join(format_sample(sample) for sample in samples)
        resume_checkpoint: dict[str, Any] | None = None
        resumed_from_checkpoint: Path | None = None
        if continue_from_checkpoint:
            resumed_from_checkpoint, resume_checkpoint = load_resume_checkpoint(resume_checkpoint_path)
            config = ModelConfig(**resume_checkpoint["config"])
            preset = str(resume_checkpoint.get("metadata", {}).get("preset") or "continued")
            stoi = resume_checkpoint["stoi"]
            itos = normalize_itos(resume_checkpoint["itos"])
            missing_chars = sorted(set(text) - set(stoi))
            if missing_chars:
                preview = "".join(missing_chars[:20])
                raise RuntimeError(
                    "Cannot continue from this checkpoint because the new dataset contains characters "
                    f"that are not in the checkpoint vocabulary: {preview!r}"
                )
            append_log(job_id, f"continue_from_checkpoint = {resumed_from_checkpoint}")
            append_log(job_id, f"checkpoint config = {config}")
        else:
            stoi, itos = build_vocab(text)
        ids = torch.tensor([stoi[ch] for ch in text], dtype=torch.long)
        if len(ids) <= config.block_size + batch_size + 2:
            raise RuntimeError(
                "Dataset is too small for the selected preset and batch size. "
                "Generate more JSONL samples or choose the tiny preset."
            )
        split = max(config.block_size + batch_size + 2, int(len(ids) * 0.92))
        split = min(split, len(ids) - config.block_size - batch_size - 2)
        train_ids = ids[:split]
        val_ids = ids[split:]
        if len(val_ids) <= config.block_size + batch_size + 2:
            val_ids = train_ids

        steps_per_epoch = max(1, math.ceil(len(train_ids) / (batch_size * config.block_size)))
        epoch_steps = math.ceil(steps_per_epoch * target_epochs) if target_epochs > 0 else fallback_max_steps
        max_steps = max(20, min(epoch_steps, 200_000))
        model = TinyCodeGPT(len(stoi), config).to(device)
        optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.1)
        if resume_checkpoint is not None:
            model.load_state_dict(resume_checkpoint["model_state_dict"])
            if resume_checkpoint.get("optimizer_state_dict"):
                optimizer.load_state_dict(resume_checkpoint["optimizer_state_dict"])
                for state in optimizer.state.values():
                    for key, value in state.items():
                        if hasattr(value, "to"):
                            state[key] = value.to(device)
        parameter_count = count_parameters(model)
        update_job(
            job_id,
            dataset_size=len(samples),
            dataset_source=dataset_source,
            dataset_path=str(dataset_path) if dataset_path else None,
            vocab_size=len(stoi),
            parameter_count=parameter_count,
            preset=preset,
            model_config=config.__dict__,
            max_steps=max_steps,
            fallback_max_steps=fallback_max_steps,
            target_epochs=target_epochs,
            steps_per_epoch=steps_per_epoch,
            resumed_from_checkpoint=str(resumed_from_checkpoint) if resumed_from_checkpoint else None,
        )
        append_log(job_id, f"dataset_size = {len(samples)} samples")
        append_log(job_id, f"dataset_source = {dataset_source}")
        if dataset_path:
            append_log(job_id, f"dataset saved: {dataset_path}")
        append_log(job_id, f"vocab_size = {len(stoi)} characters")
        append_log(job_id, f"parameters = {parameter_count:,}")
        append_log(
            job_id,
            f"target_epochs = {target_epochs:g}, steps_per_epoch ≈ {steps_per_epoch}, "
            f"planned_steps = {max_steps}, fallback_max_steps = {fallback_max_steps}",
        )

        log_every = max(10, max_steps // 18)
        start_time = time.time()
        for step in range(1, max_steps + 1):
            if stop_event.is_set():
                update_job(job_id, status="stopped", finished_at=utc_now())
                append_log(job_id, "training stopped by user")
                return
            x, y = make_batch(train_ids, config.block_size, batch_size, device)
            _, loss = model(x, y)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()

            if step == 1 or step % log_every == 0 or step == max_steps:
                val_loss = evaluate(model, val_ids, config.block_size, min(batch_size, 32), device)
                elapsed = max(time.time() - start_time, 1e-6)
                tokens_per_second = int(step * batch_size * config.block_size / elapsed)
                metric = {
                    "step": step,
                    "train_loss": round(float(loss.item()), 6),
                    "val_loss": round(float(val_loss), 6),
                    "tokens_per_second": tokens_per_second,
                    "epoch": round(step / steps_per_epoch, 4),
                }
                append_metric(job_id, metric)
                append_log(
                    job_id,
                    f"step {step:5d}/{max_steps} | epoch {metric['epoch']:.2f}/{target_epochs:g} | train_loss {metric['train_loss']:.4f} | val_loss {metric['val_loss']:.4f} | {tokens_per_second:,} tok/s",
                )

        checkpoint_path = CHECKPOINT_DIR / f"{job_id}.pt"
        torch.save(
            {
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "config": config.__dict__,
                "stoi": stoi,
                "itos": itos,
                "samples": samples[:120],
                "metadata": {
                    "preset": preset,
                    "model_config": config.__dict__,
                    "samples_per_category": samples_per_category,
                    "max_steps": max_steps,
                    "fallback_max_steps": fallback_max_steps,
                    "target_epochs": target_epochs,
                    "steps_per_epoch": steps_per_epoch,
                    "batch_size": batch_size,
                    "learning_rate": learning_rate,
                    "dataset_source": dataset_source,
                    "dataset_path": str(dataset_path) if dataset_path else None,
                    "dataset_size": len(samples),
                    "resumed_from_checkpoint": str(resumed_from_checkpoint) if resumed_from_checkpoint else None,
                    "created_at": utc_now(),
                },
            },
            checkpoint_path,
        )
        update_job(job_id, status="completed", checkpoint_path=str(checkpoint_path), finished_at=utc_now())
        append_log(job_id, f"checkpoint saved: {checkpoint_path}")
    except Exception as error:
        update_job(job_id, status="failed", error=str(error), finished_at=utc_now())
        append_log(job_id, traceback.format_exc())
    finally:
        with LOCK:
            if ACTIVE_JOB_ID == job_id:
                ACTIVE_JOB_ID = None


def latest_checkpoint() -> Path | None:
    checkpoints = sorted(CHECKPOINT_DIR.glob("*.pt"), key=lambda item: item.stat().st_mtime, reverse=True)
    return checkpoints[0] if checkpoints else None


def load_checkpoint(path_value: str | None):
    if torch is None:
        raise RuntimeError("PyTorch is not installed in this Python environment.")
    checkpoint_path = Path(path_value) if path_value else latest_checkpoint()
    if checkpoint_path is None or not checkpoint_path.exists():
        raise RuntimeError("No TinyCodeGPT checkpoint found. Train a model first.")
    checkpoint = torch.load(checkpoint_path, map_location="cpu")
    config = ModelConfig(**checkpoint["config"])
    stoi = checkpoint["stoi"]
    itos = normalize_itos(checkpoint["itos"])
    model = TinyCodeGPT(len(stoi), config)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()
    return checkpoint_path, model, config, stoi, itos


def extract_code(text: str) -> str:
    marker = "<code>\n"
    start = text.find(marker)
    if start >= 0:
        text = text[start + len(marker) :]
    end = text.find("</code>")
    if end >= 0:
        text = text[:end]
    return text.strip()


def encode_png_artifact(path: Path) -> dict[str, str]:
    payload = base64.b64encode(path.read_bytes()).decode("ascii")
    return {
        "filename": path.name,
        "path": str(path),
        "mime_type": "image/png",
        "data_url": f"data:image/png;base64,{payload}",
    }


def run_generated_code(code: str, timeout_seconds: int = 12) -> dict[str, Any]:
    if not code.strip():
        return {"output": [], "error": "Model returned empty code.", "artifacts": []}
    try:
        assert_safe_sample_code(code, 1)
    except ValueError as error:
        return {"output": [], "error": str(error), "artifacts": []}
    run_id = uuid.uuid4().hex[:10]
    script_path = RUN_OUTPUT_DIR / f"generated_{run_id}.py"
    capture_header = """
# AlgoLab runs matplotlib with a non-interactive backend and captures figures
# after the generated code finishes. Ignore the expected plt.show() warning.
import warnings as _algolab_warnings
_algolab_warnings.filterwarnings(
    "ignore",
    message="FigureCanvasAgg is non-interactive.*",
    category=UserWarning,
)
"""
    capture_footer = f"""

# AlgoLab captures matplotlib figures after generated code runs.
try:
    import os as _algolab_os
    import sys as _algolab_sys
    try:
        import matplotlib.pyplot as _algolab_plt
    except Exception:
        _algolab_plt = None
    _algolab_plot_dir = r"{RUN_OUTPUT_DIR}"
    if _algolab_plt is not None:
        for _algolab_index, _algolab_num in enumerate(_algolab_plt.get_fignums(), start=1):
            _algolab_figure = _algolab_plt.figure(_algolab_num)
            _algolab_path = _algolab_os.path.join(_algolab_plot_dir, "{run_id}_plot_" + str(_algolab_index) + ".png")
            _algolab_figure.savefig(_algolab_path, bbox_inches="tight")
except Exception as _algolab_plot_error:
    print("[AlgoLab plot capture failed] " + str(_algolab_plot_error), file=_algolab_sys.stderr)
"""
    script_path.write_text(f"{capture_header}\n{code.rstrip()}\n{capture_footer}", encoding="utf-8")
    env = {
        **os.environ,
        "PYTHONUTF8": "1",
        "PYTHONIOENCODING": "utf-8",
        "MPLBACKEND": "Agg",
    }
    try:
        completed = subprocess.run(
            [sys.executable, "-X", "utf8", str(script_path)],
            cwd=str(RUN_OUTPUT_DIR),
            env=env,
            text=True,
            capture_output=True,
            timeout=timeout_seconds,
            check=False,
        )
    except subprocess.TimeoutExpired:
        return {"output": [], "error": f"Generated code timed out after {timeout_seconds} seconds.", "artifacts": []}
    output = (completed.stdout or "").strip().splitlines()
    stderr_lines = (completed.stderr or "").strip().splitlines()
    filtered_stderr_lines: list[str] = []
    skip_show_line = False
    for line in stderr_lines:
        if "FigureCanvasAgg is non-interactive" in line:
            skip_show_line = True
            continue
        if skip_show_line and line.strip().startswith("plt.show("):
            skip_show_line = False
            continue
        skip_show_line = False
        filtered_stderr_lines.append(line)
    stderr = "\n".join(filtered_stderr_lines).strip()
    error = stderr or None
    if completed.returncode != 0 and not error:
        error = f"Generated code exited with status {completed.returncode}."
    artifacts = [encode_png_artifact(path) for path in sorted(RUN_OUTPUT_DIR.glob(f"{run_id}_plot_*.png"))]
    return {
        "output": output[-80:],
        "error": error,
        "artifacts": artifacts,
        "return_code": completed.returncode,
        "script_path": str(script_path),
    }


def generate_code(request: dict[str, Any]) -> dict[str, Any]:
    checkpoint_path, model, config, stoi, itos = load_checkpoint(request.get("checkpoint_path"))
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    task = str(request.get("task") or "").strip()
    if not task:
        raise RuntimeError("Task cannot be empty.")
    prompt = f"<task>\n{task}\n</task>\n<code>\n"
    encoded = [stoi[ch] for ch in prompt if ch in stoi]
    if not encoded:
        raise RuntimeError("Prompt contains no known characters from the tokenizer vocabulary.")
    idx = torch.tensor([encoded], dtype=torch.long, device=device)
    max_new_tokens = max(80, min(int(request.get("max_new_tokens") or 900), 2000))
    temperature = max(0.2, min(float(request.get("temperature") or 0.75), 1.5))
    stop_sequence = [stoi[ch] for ch in "\n</code>" if ch in stoi]
    with torch.no_grad():
        generated = model.generate(
            idx,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            stop_sequence=stop_sequence if stop_sequence else None,
        )
    text = "".join(itos[int(token)] for token in generated[0].detach().cpu().tolist())
    code = extract_code(text)
    return {
        "task": task,
        "code": code,
        "truncated": "</code>" not in text,
        "checkpoint_path": str(checkpoint_path),
    }


def execute_code(request: dict[str, Any]) -> dict[str, Any]:
    timeout_seconds = max(1, min(int(request.get("timeout_seconds") or 12), 30))
    code = str(request.get("code") or "").strip()
    result = run_generated_code(code, timeout_seconds=timeout_seconds)
    return {"code": code, **result}


def status_payload() -> dict[str, Any]:
    checkpoints = sorted(CHECKPOINT_DIR.glob("*.pt"), key=lambda item: item.stat().st_mtime, reverse=True)
    torch_available = torch is not None
    cuda_available = bool(torch_available and torch.cuda.is_available())
    return {
        "ok": True,
        "python": sys.version.split()[0],
        "torch_available": torch_available,
        "torch_version": getattr(torch, "__version__", None) if torch_available else None,
        "cuda_available": cuda_available,
        "cuda_device": torch.cuda.get_device_name(0) if cuda_available else None,
        "device": "cuda" if cuda_available else "cpu",
        "working_dir": str(RUN_DIR),
        "active_job_id": ACTIVE_JOB_ID,
        "checkpoints": [str(item) for item in checkpoints[:10]],
    }


class TinyCodeGPTHandler(BaseHTTPRequestHandler):
    def log_message(self, format: str, *args: Any) -> None:
        print(f"[tinycodegpt-runner] {self.address_string()} - {format % args}")

    def do_OPTIONS(self) -> None:
        respond(self, 200, {"ok": True})

    def do_GET(self) -> None:
        parsed = urlparse(self.path)
        path = parsed.path.rstrip("/") or "/"
        if path in {"/", "/status", "/health"}:
            respond(self, 200, status_payload())
            return
        if path.startswith("/jobs/"):
            job_id = path.split("/")[-1]
            with LOCK:
                job = JOBS.get(job_id)
            if not job:
                respond(self, 404, {"message": "Job not found."})
                return
            respond(self, 200, job)
            return
        respond(self, 404, {"message": "Unknown route."})

    def do_POST(self) -> None:
        global ACTIVE_JOB_ID
        parsed = urlparse(self.path)
        path = parsed.path.rstrip("/") or "/"
        try:
            if path == "/train":
                body = read_body(self)
                with LOCK:
                    if ACTIVE_JOB_ID and JOBS.get(ACTIVE_JOB_ID, {}).get("status") in {"queued", "running"}:
                        respond(self, 409, {"message": f"Training job already running: {ACTIVE_JOB_ID}"})
                        return
                    job_id = uuid.uuid4().hex[:12]
                    job = {
                        "id": job_id,
                        "status": "queued",
                        "created_at": utc_now(),
                        "logs": [],
                        "metrics": [],
                    }
                    JOBS[job_id] = job
                    STOP_EVENTS[job_id] = threading.Event()
                    ACTIVE_JOB_ID = job_id
                thread = threading.Thread(target=run_training, args=(job_id, body), daemon=True)
                thread.start()
                respond(self, 200, job)
                return

            if path.startswith("/jobs/") and path.endswith("/stop"):
                parts = path.split("/")
                job_id = parts[2] if len(parts) >= 3 else ""
                if job_id in STOP_EVENTS:
                    STOP_EVENTS[job_id].set()
                with LOCK:
                    job = JOBS.get(job_id)
                if not job:
                    respond(self, 404, {"message": "Job not found."})
                    return
                respond(self, 200, job)
                return

            if path == "/generate":
                body = read_body(self)
                result = generate_code(body)
                respond(self, 200, result)
                return

            if path == "/execute":
                body = read_body(self)
                result = execute_code(body)
                respond(self, 200, result)
                return

            respond(self, 404, {"message": "Unknown route."})
        except Exception as error:
            respond(self, 500, {"message": str(error), "traceback": traceback.format_exc()})


def main() -> None:
    parser = argparse.ArgumentParser(description="Local TinyCodeGPT training runner for AlgoLab.")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=4877)
    args = parser.parse_args()
    server = ThreadingHTTPServer((args.host, args.port), TinyCodeGPTHandler)
    print(f"TinyCodeGPT local runner listening on http://{args.host}:{args.port}")
    print(f"Working directory: {RUN_DIR}")
    print("Press Ctrl+C to stop.")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nStopping TinyCodeGPT local runner.")
    finally:
        server.server_close()


if __name__ == "__main__":
    main()