runapprentice 0.4.3__tar.gz → 0.4.4__tar.gz

{runapprentice-0.4.3 → runapprentice-0.4.4}/.gitignore

@@ -53,3 +53,8 @@ docs/personal_learning.md
 spike/Capture-*.png
 apprentice-ui/ClaudeDesign/screenshots/localhost-*.png
 supabase-magic-link-template.html
+
+# Remotion scratch: frame exports for review, retired captures, render output.
+# 30MB of JPEG/PNG working files — never source.
+apprentice-video/tmp/
+apprentice-video/out/

{runapprentice-0.4.3 → runapprentice-0.4.4}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: runapprentice
-Version: 0.4.3
+Version: 0.4.4
 Summary: Cut LLM cost without losing quality: capture traffic, build golden datasets, and optimize prompts.
 Project-URL: Homepage, https://runapprentice.com
 Project-URL: Documentation, https://docs.runapprentice.com
@@ -22,7 +22,7 @@ Requires-Dist: pydantic>=2.0
 Provides-Extra: langchain
 Requires-Dist: langchain-core>=0.3; extra == 'langchain'
 Provides-Extra: local
-Requires-Dist: mlx-lm<0.32,>=0.31.3; extra == 'local'
+Requires-Dist: mlx-lm<0.33,>=0.32.2; extra == 'local'
 Provides-Extra: optimize
 Requires-Dist: dspy<4,>=3.2.1; extra == 'optimize'
 Description-Content-Type: text/markdown

{runapprentice-0.4.3 → runapprentice-0.4.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "runapprentice"
-version = "0.4.3"
+version = "0.4.4"
 description = "Cut LLM cost without losing quality: capture traffic, build golden datasets, and optimize prompts."
 readme = "README.md"
 requires-python = ">=3.11"
@@ -32,7 +32,7 @@ Documentation = "https://docs.runapprentice.com"
 
 [project.optional-dependencies]
 langchain = ["langchain-core>=0.3"]
-local = ["mlx-lm>=0.31.3,<0.32"]
+local = ["mlx-lm>=0.32.2,<0.33"]
 optimize = ["dspy>=3.2.1,<4"]
 
 [dependency-groups]
@@ -57,6 +57,9 @@ only-include = ["src/runapprentice", "tests", "README.md", "LICENSE", "pyproject
 line-length = 100
 src = ["src", "tests"]
 
+[tool.ruff.lint]
+select = ["E4", "E7", "E9", "F"]
+
 [tool.mypy]
 strict = true
 mypy_path = "src"

runapprentice-0.4.4/src/runapprentice/_local_ops.py

@@ -0,0 +1,171 @@
+"""Process ownership and durable local training status. No MLX imports."""
+
+from __future__ import annotations
+
+import datetime as dt
+import fcntl
+import json
+import os
+import signal
+import tempfile
+import threading
+import time
+import uuid
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any, Iterator
+
+from runapprentice._train_output import status_command, watch_command
+from runapprentice.client import ApprenticeError
+
+_owner: dict[str, Any] | None = None
+_thread: int | None = None
+_fd: int | None = None
+json_progress = False
+
+
+def lock_fd() -> int | None:
+    return _fd
+
+
+# ponytail: the lock lives in this user's home, so it serializes runs per macOS account. Two
+# accounts training at once can still compete for the GPU; a shared lock would need a
+# world-writable path and its permission problems. Docs say "per user account".
+def _root() -> Path:
+    return Path.home() / ".apprentice" / "local-runs"
+
+
+def _read(path: Path) -> dict[str, Any]:
+    try:
+        value = json.loads(path.read_text(encoding="utf-8"))
+        return value if isinstance(value, dict) else {}
+    except FileNotFoundError:
+        return {}
+
+
+def _write(path: Path, value: dict[str, Any]) -> None:
+    name: str | None = None
+    try:
+        with tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", dir=path.parent,
+                                         delete=False) as out:
+            name = out.name
+            json.dump(value, out)
+            out.write("\n")
+        os.replace(name, path)
+    finally:
+        if name is not None:
+            Path(name).unlink(missing_ok=True)
+
+
+def update(**values: Any) -> None:
+    if _owner is not None:
+        if values.get("state") in {"failed", "completed"}:
+            values.setdefault("finished", dt.datetime.now(dt.UTC).isoformat())
+        _owner.update(values)
+        _write(Path(_owner["status_path"]), _owner)
+
+
+def snapshot() -> dict[str, Any]:
+    return dict(_owner) if _owner is not None else {}
+
+
+def status(task: str) -> dict[str, Any]:
+    from runapprentice.local import _local_run_dir
+
+    path = _local_run_dir(task) / "status.json"
+    value = _read(path)
+    if not value:
+        return {"task": task, "state": "none"}
+    if value.get("state") == "running":
+        try:
+            with (_root() / ".lock").open("r+") as lock:
+                try:
+                    fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                    value["state"] = "stopped"
+                except BlockingIOError:
+                    holder = _read(_root() / "holder.json")
+                    if holder.get("run_id") != value.get("run_id"):
+                        value["state"] = "stopped"
+        except FileNotFoundError:
+            value["state"] = "stopped"
+    return value
+
+
+def _interrupt(signum: int, frame: Any) -> None:
+    reason = ("the terminal closed" if signum == signal.SIGHUP
+              else "the process was told to quit, often by a command timeout")
+    raise KeyboardInterrupt(f"Stopped because {reason}.")
+
+
+def _acquire(lock: Any) -> None:
+    # status() probes this lock for an instant, so one refusal may only mean a poll was
+    # looking. A real run holds it for much longer than these retries.
+    for attempt in range(3):
+        try:
+            fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
+            return
+        except BlockingIOError:
+            if attempt == 2:
+                raise
+            time.sleep(0.05)
+
+
+@contextmanager
+def training(task: str) -> Iterator[None]:
+    global _owner, _thread, _fd
+    from runapprentice.local import _local_run_dir
+
+    folder = _local_run_dir(task)
+    if _owner is not None and _thread == threading.get_ident():
+        if _owner["task"] != task:
+            raise ApprenticeError("Another local task is already running.", code="busy")
+        yield
+        return
+    _root().mkdir(parents=True, exist_ok=True)
+    with (_root() / ".lock").open("a+") as lock:
+        try:
+            _acquire(lock)
+        except BlockingIOError as exc:
+            holder = _read(_root() / "holder.json")
+            command = status_command(str(holder.get("task", task)))
+            raise ApprenticeError(
+                f"Local training is busy: task {holder.get('task', 'unknown')}, "
+                f"pid {holder.get('pid', 'unknown')}. Check: {command}", code="busy",
+                remediation=command,
+            ) from exc
+        previous = {}
+        try:
+            folder.mkdir(parents=True, exist_ok=True)
+            _fd = lock.fileno()
+            _thread = threading.get_ident()
+            _owner = {"task": task, "pid": os.getpid(), "run_id": uuid.uuid4().hex,
+                      "state": "running", "phase": "preflight", "step": None,
+                      "total": None, "percent": None, "time_left": None,
+                      "train_loss": None, "val_loss": None, "val_step": None,
+                      "train_points": [], "val_points": [], "memory_points": [],
+                      "speed_points": [], "tokens_per_second": None,
+                      "peak_memory": None,
+                      "effort": None, "started": dt.datetime.now(dt.UTC).isoformat(),
+                      "finished": None, "status_path": str(folder / "status.json")}
+            _owner.update(log_path=str(folder / "train.log"),
+                          status_command=status_command(task), watch_command=watch_command(task))
+            _write(_root() / "holder.json", _owner)
+            update()
+            with (folder / "train.log").open("w", encoding="utf-8") as log:
+                log.write(f"Training {task!r}, pid {os.getpid()}, started {_owner['started']}\n")
+            if threading.current_thread() is threading.main_thread():
+                for sig in (signal.SIGTERM, signal.SIGHUP):
+                    previous[sig] = signal.signal(sig, _interrupt)
+            yield
+        except BaseException as exc:
+            try:
+                update(state="failed", message=str(exc) or "Local training was interrupted.",
+                       finished=dt.datetime.now(dt.UTC).isoformat())
+            except OSError:
+                # ponytail: preserve the original failure if the disk cannot save status.
+                pass
+            raise
+        finally:
+            for sig, handler in previous.items():
+                signal.signal(sig, handler)
+            _owner, _thread, _fd = None, None, None

runapprentice-0.4.4/src/runapprentice/_mlx_eval.py

@@ -0,0 +1,70 @@
+"""Generate held-out answers with one model load: `python -m runapprentice._mlx_eval`.
+
+Reads one JSON string (the prompt) per line from --prompts and writes one JSON
+string (the answer) per line to --out. `mlx_lm generate` per row reloaded the
+model every time. The memory guard from _mlx_lora covers this step too.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+from typing import Any, cast
+
+import mlx.core as mx
+
+from runapprentice._mlx_lora import memory_ceiling_gb
+from runapprentice.local import MEMORY_GUARD_EXIT
+
+# Generation stops at the end of the answer; this only caps a model that never stops.
+# mlx_lm generate's default of 100 tokens cut off real JSON answers mid-object.
+MAX_TOKENS = 1024
+
+
+def main() -> None:
+    from mlx_lm import generate, load
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", required=True)
+    parser.add_argument("--adapter-path", required=True)
+    parser.add_argument("--prompts", type=Path, required=True)
+    parser.add_argument("--out", type=Path, required=True)
+    args = parser.parse_args()
+
+    ceiling = memory_ceiling_gb()
+    loaded = load(args.model, adapter_path=args.adapter_path)
+    # Any: mlx-lm's tokenizer wrapper forwards apply_chat_template untyped.
+    model, tokenizer = loaded[0], cast(Any, loaded[1])
+    with args.prompts.open(encoding="utf-8") as src, args.out.open("w", encoding="utf-8") as out:
+        for line in src:
+            prompt = tokenizer.apply_chat_template(
+                [{"role": "user", "content": json.loads(line)}],
+                tokenize=False,
+                add_generation_prompt=True,
+                # Qwen3.5 reasons in free text by default; the scorer needs the answer.
+                enable_thinking=False,
+            )
+            answer = generate(
+                model,
+                tokenizer,
+                tokenizer.encode(prompt, add_special_tokens=False),
+                max_tokens=MAX_TOKENS,
+            )
+            out.write(json.dumps(answer, ensure_ascii=False) + "\n")
+            out.flush()
+            if os.environ.get("APPRENTICE_MLX_PROGRESS"):
+                print('APPRENTICE_PROGRESS {"row": 1}', flush=True)
+            peak = mx.get_peak_memory() / 1e9
+            if peak > ceiling:
+                print(
+                    f"[apprentice] memory guard: eval peaked at {peak:.1f} GB, over the "
+                    f"{ceiling:.1f} GB safe limit for this Mac's GPU. Stopping.",
+                    flush=True,
+                )
+                raise SystemExit(MEMORY_GUARD_EXIT)
+
+
+if __name__ == "__main__":
+    main()

runapprentice-0.4.4/src/runapprentice/_mlx_lora.py

@@ -0,0 +1,243 @@
+"""`python -m mlx_lm lora` with two changes for long-row Qwen3.5 training.
+
+1. Speed and memory. In training mode mlx-lm runs Qwen3.5's linear-attention
+   recurrence (`gated_delta_ops`) one token at a time, keeping a 2 MB float32
+   state per token per layer for backprop. An 800-token row took 26-125 s per
+   step and ~23 GB. `gated_delta_chunked` computes the same recurrence 64
+   tokens at a time with matmuls (the chunkwise "WY" form of the delta rule),
+   matching the reference to float32 rounding. Model quality is unchanged.
+2. A memory guard. MLX's own memory limit is a guideline, not a cap, so a run
+   that outgrows the GPU working set pushes the Mac into swap. The guard stops
+   training with exit code MEMORY_GUARD_EXIT as soon as a step's peak passes
+   the ceiling, before the next, possibly longer, step can go further.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+import sys
+from pathlib import Path
+from typing import Any, Callable
+
+import mlx.core as mx
+
+from runapprentice.local import MEMORY_GUARD_EXIT, _check_token_lengths
+from runapprentice.client import ApprenticeError
+# Fraction of Metal's recommended working set a step may peak at. The rest is
+# headroom for a longer row than the one that set the peak.
+MEMORY_GUARD_FRACTION = 0.85
+CHUNK = 64
+
+
+def _unit_lower_inverse(n: mx.array) -> mx.array:
+    # Forward substitution for (I + N)^-1, N strictly lower triangular. Row i is
+    # e_i - N[i, :i] @ rows[:i]. A power series ((I - N)(I + N^2)...) is shorter
+    # but cancels catastrophically in float32 when keys are correlated, as in text.
+    c = n.shape[-1]
+    eye = mx.eye(c, dtype=n.dtype)
+    rows = [mx.broadcast_to(eye[0], n.shape[:-2] + (c,))]
+    for i in range(1, c):
+        prev = mx.stack(rows, axis=-2)  # [..., i, C]
+        rows.append(eye[i] - (n[..., i : i + 1, :i] @ prev)[..., 0, :])
+    return mx.stack(rows, axis=-2)
+
+
+def gated_delta_chunked(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: mx.array,
+    chunk: int = CHUNK,
+) -> tuple[mx.array, mx.array]:
+    """Same result as mlx_lm `gated_delta_ops` for a scalar gate and no mask.
+
+    q, k: [B, T, Hk, Dk]; v: [B, T, Hv, Dv]; g, beta: [B, T, Hv]; state: [B, Hv, Dv, Dk].
+    Per token the reference does S = g*S; d = beta*(v - S k); S += d k^T; y = S q.
+    """
+    B, T, Hk, _ = q.shape
+    Hv = v.shape[-2]
+    if (repeat := Hv // Hk) > 1:
+        q = mx.repeat(q, repeat, -2)
+        k = mx.repeat(k, repeat, -2)
+    in_dtype = q.dtype
+    q, k, v = (x.astype(mx.float32).transpose(0, 2, 1, 3) for x in (q, k, v))
+    # Clamp before log: log's backward divides by its input, even with zero cotangent.
+    log_g = mx.log(mx.maximum(g.astype(mx.float32), 1e-30)).transpose(0, 2, 1)
+    beta = beta.astype(mx.float32).transpose(0, 2, 1)
+    # Padding tokens have k = 0, beta = 0 and g = 1, so they leave the state untouched.
+    pad = (-T) % chunk
+    if pad:
+        q, k, v = (mx.pad(x, [(0, 0), (0, 0), (0, pad), (0, 0)]) for x in (q, k, v))
+        log_g = mx.pad(log_g, [(0, 0), (0, 0), (0, pad)])
+        beta = mx.pad(beta, [(0, 0), (0, 0), (0, pad)])
+    n = (T + pad) // chunk
+    q, k, v = (x.reshape(B, Hv, n, chunk, x.shape[-1]) for x in (q, k, v))
+    log_g = log_g.reshape(B, Hv, n, chunk)
+    beta = beta.reshape(B, Hv, n, chunk)
+
+    cum = mx.cumsum(log_g, axis=-1)  # log of the decay from chunk start to token i
+    idx = mx.arange(chunk)
+    strict = idx[:, None] > idx[None, :]
+    inclusive = idx[:, None] >= idx[None, :]
+    diff = cum[..., :, None] - cum[..., None, :]
+    decay_strict = mx.where(strict, mx.exp(mx.where(strict, diff, 0.0)), 0.0)
+    decay_inclusive = mx.where(inclusive, mx.exp(mx.where(inclusive, diff, 0.0)), 0.0)
+    # Within a chunk the deltas solve (I + A) D = beta * (V - Gamma K S0^T).
+    a = beta[..., None] * decay_strict * (k @ k.swapaxes(-1, -2))
+    a_inv = _unit_lower_inverse(a)
+    gamma = mx.exp(cum)[..., None]
+    w = a_inv @ (beta[..., None] * gamma * k)
+    u = a_inv @ (beta[..., None] * v)
+    attn = decay_inclusive * (q @ k.swapaxes(-1, -2))
+    to_end = mx.exp(cum[..., -1:] - cum)[..., None]
+
+    ys = []
+    for c in range(n):
+        state_t = state.swapaxes(-1, -2)
+        delta = u[:, :, c] - w[:, :, c] @ state_t
+        ys.append(gamma[:, :, c] * (q[:, :, c] @ state_t) + attn[:, :, c] @ delta)
+        state = mx.exp(cum[:, :, c, -1])[..., None, None] * state + (
+            delta.swapaxes(-1, -2) @ (to_end[:, :, c] * k[:, :, c])
+        )
+    y = mx.concatenate(ys, axis=2)[:, :, :T].transpose(0, 2, 1, 3)
+    return y.astype(in_dtype), state
+
+
+def _install_chunked_recurrence() -> None:
+    from mlx_lm.models import gated_delta
+
+    reference: Callable[..., tuple[mx.array, mx.array]] = gated_delta.gated_delta_ops
+
+    def gated_delta_ops(
+        q: mx.array,
+        k: mx.array,
+        v: mx.array,
+        g: mx.array,
+        beta: mx.array,
+        state: mx.array | None = None,
+        mask: mx.array | None = None,
+    ) -> tuple[mx.array, mx.array]:
+        # Training passes no mask and a per-head scalar gate; anything else keeps the reference.
+        if mask is not None or g.ndim != 3:
+            return reference(q, k, v, g, beta, state, mask)
+        if state is None:
+            B, _, _, Dk = q.shape
+            Hv, Dv = v.shape[-2:]
+            state = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32)
+        return gated_delta_chunked(q, k, v, g, beta, state)
+
+    # gated_delta_update looks this name up in its module at call time.
+    gated_delta.gated_delta_ops = gated_delta_ops
+
+
+def memory_ceiling_gb() -> float:
+    return float(mx.device_info()["max_recommended_working_set_size"]) * MEMORY_GUARD_FRACTION / 1e9
+
+
+def _install_memory_guard() -> None:
+    from mlx_lm import lora
+    from mlx_lm.tuner.callbacks import get_reporting_callbacks as original
+
+    ceiling = memory_ceiling_gb()
+    # local.py reads the last report to estimate the full run's time.
+    stats_path = os.environ.get("APPRENTICE_MLX_STATS")
+
+    # Duck-typed: the trainer only calls these two methods.
+    class MemoryGuard:
+        def __init__(self, inner: Any) -> None:
+            self.inner = inner
+
+        def on_train_loss_report(self, train_info: dict[str, Any]) -> None:
+            if not math.isfinite(float(train_info["train_loss"])):
+                print("Training stopped because the training loss is NaN or infinite.", flush=True)
+                raise SystemExit(87)
+            if os.environ.get("APPRENTICE_MLX_PROGRESS"):
+                progress = {
+                    key: value for key, value in train_info.items()
+                    if key == "iteration" or (
+                        key in {"train_loss", "peak_memory", "tokens_per_second"}
+                        and math.isfinite(float(value))
+                    )
+                }
+                print(
+                    "APPRENTICE_PROGRESS "
+                    + json.dumps(progress, default=float, allow_nan=False),
+                    flush=True,
+                )
+            if self.inner is not None:
+                self.inner.on_train_loss_report(train_info)
+            if stats_path:
+                Path(stats_path).write_text(json.dumps(train_info, default=float), encoding="utf-8")
+            # local.py requests a report every step, before another batch runs.
+            peak = float(train_info["peak_memory"])  # GB, 1e9 bytes, as mlx-lm reports it
+            if peak > ceiling:
+                print(
+                    f"[apprentice] memory guard: training peaked at {peak:.1f} GB, over the "
+                    f"{ceiling:.1f} GB safe limit for this Mac's GPU. Stopping.",
+                    flush=True,
+                )
+                raise SystemExit(MEMORY_GUARD_EXIT)
+
+        def on_val_loss_report(self, val_info: dict[str, Any]) -> None:
+            if (
+                os.environ.get("APPRENTICE_MLX_PROGRESS")
+                and math.isfinite(float(val_info["val_loss"]))
+            ):
+                print(
+                    "APPRENTICE_PROGRESS "
+                    + json.dumps(
+                        {key: val_info[key] for key in ("iteration", "val_loss")},
+                        default=float,
+                        allow_nan=False,
+                    ),
+                    flush=True,
+                )
+            if self.inner is not None:
+                self.inner.on_val_loss_report(val_info)
+
+    def get_reporting_callbacks(*args: Any, **kwargs: Any) -> MemoryGuard:
+        return MemoryGuard(original(*args, **kwargs))
+
+    # lora.run looks this name up in its own module at call time.
+    setattr(lora, "get_reporting_callbacks", get_reporting_callbacks)
+
+
+def main() -> None:
+    _install_chunked_recurrence()
+    _install_memory_guard()
+    from mlx_lm import lora
+
+    from mlx_lm.tuner.datasets import create_dataset
+
+    original_load = getattr(lora, "load_dataset")
+
+    def load_largest(args: Any, tokenizer: Any) -> Any:
+        train, valid, test = original_load(args, tokenizer)
+        for name, dataset in (("training", train), ("held-out", valid)):
+            try:
+                _check_token_lengths(dataset, args.max_seq_length, name)
+            except ApprenticeError as exc:
+                print(str(exc), flush=True)
+                raise SystemExit(86)
+        if not os.environ.get("APPRENTICE_MLX_STATS"):
+            return train, valid, test
+        # Reuse the training tokenizer and chat template, without loading another model.
+        indices = sorted(
+            range(len(train)),
+            key=lambda i: min(len(train.process(train[i])[0]), args.max_seq_length),
+        )[-args.batch_size :]
+        largest = create_dataset([train[i] for i in indices], tokenizer, args)
+        return largest, largest, test
+
+    setattr(lora, "load_dataset", load_largest)
+
+    sys.argv[0] = "mlx_lm.lora"
+    getattr(lora, "main")()  # untyped in mlx-lm
+
+
+if __name__ == "__main__":
+    main()