blitz-cli 0.5.0__tar.gz → 0.6.0__tar.gz

{blitz_cli-0.5.0 → blitz_cli-0.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: blitz-cli
-Version: 0.5.0
+Version: 0.6.0
 Summary: Developer CLI for Blitz: local-first tracing/init, codebase instrumentation scan, and scaffolds for both a QLoRA workflow trainer and a coding-agent session-to-fine-tuning-corpus watcher.
 License: MIT
 License-File: LICENSE

{blitz_cli-0.5.0 → blitz_cli-0.6.0}/blitz_cli/__init__.py

@@ -30,4 +30,4 @@ Commands:
 
 __all__ = ["__version__"]
 
-__version__ = "0.5.0"
+__version__ = "0.6.0"

{blitz_cli-0.5.0 → blitz_cli-0.6.0}/blitz_cli/_scaffold.py

@@ -23,8 +23,17 @@ from string import Template
 
 def clamp_seq_len(raw: int) -> int:
     """Clamp a raw sequence length (p95_input + max_output) to the trainer's
-    supported range. Defaults to 2048 when nothing usable is supplied."""
-    return max(1024, min(8192, int(raw or 0) or 2048))
+    supported range. Defaults to 2048 when nothing usable is supplied.
+
+    Upper bound is 16384, not 8192 -- measured (not estimated) against a
+    real code-assist corpus + the real Qwen tokenizer with a full
+    tool-schema block included on every example: at 8192 the schema tax
+    alone costs +4.4pp over-budget turns versus its own no-schema
+    baseline, but at 16384 that gap shrinks to +0.8pp (nearly free
+    relative to baseline at that length). This only permits more headroom
+    for callers who ask for it -- any caller requesting <=8192 behaves
+    identically to before."""
+    return max(1024, min(16384, int(raw or 0) or 2048))
 
 
 def build_config(
@@ -66,6 +75,27 @@ def build_config(
     }
 
 
+def apply_batch_override(cfg: dict, per_device_batch_size: int) -> dict:
+    """Override just the per-device batch size from build_config's
+    params_b-driven bucket, independent of LoRA rank.
+
+    ``params_b`` ties batch size and LoRA rank together as one bucket, but
+    they hit different memory limits (batch size scales per-step activation/
+    logits memory; LoRA rank scales adapter memory) -- a card that's tight on
+    one but fine on the other has no way to ask for just one via ``params_b``
+    alone. Rescales ``grad_accum`` to hold the effective batch size
+    (``per_device_batch_size * grad_accum``) steady, so unlike ``seq_len``
+    this has no training-dynamics tradeoff: it only trades step count for
+    peak memory.
+    """
+    effective_batch = cfg["per_device_batch_size"] * cfg["grad_accum"]
+    return {
+        **cfg,
+        "per_device_batch_size": per_device_batch_size,
+        "grad_accum": max(1, round(effective_batch / per_device_batch_size)),
+    }
+
+
 def _template(name: str) -> str:
     return resources.files("blitz_cli.templates").joinpath(name).read_text(
         encoding="utf-8"
@@ -109,7 +139,7 @@ def render(out_dir: Path, workflow: str) -> None:
         _write(out_dir / name, Template(_template(name + ".tmpl")).safe_substitute(subs))
 
 
-_DEFAULT_CODE_ASSIST_BASE_MODEL = "Qwen/Qwen2.5-Coder-14B-Instruct"
+_DEFAULT_CODE_ASSIST_BASE_MODEL = "Qwen/Qwen2.5-Coder-7B-Instruct"
 
 
 def render_code_assist(

blitz_cli-0.6.0/blitz_cli/_tool_schemas.py

@@ -0,0 +1,189 @@
+"""Formal JSON-Schema tool definitions for training-time `tools=` inclusion
+(see `codeassist_train.py.tmpl`'s `to_text()`).
+
+Without this, the model never sees a formal tool schema during training —
+it only ever learns bare observed `(name, arguments)` pairs shaped like
+this environment's own tool vocabulary. That risks teaching narrow
+memorization of one fixed vocabulary rather than the general "read a given
+schema, comply with it" skill that lets a model like Claude drive arbitrary
+harnesses (a harness sends its own tool schemas at inference time; the
+model reads whatever it's handed). Training *with* `tools=` included is
+meant to teach that general skill instead.
+
+Field shapes here are sampled directly from real `tool_input` values in
+this environment's own corpus (`data/activity.db`) and from
+`_activity.py`'s own summarizer/extractor functions for the 11 vanilla
+Claude Code tools -- not guessed. `required` is a schema-design judgment
+call layered on top in every case: `_activity.py` itself never enforces
+required-ness (every reader is a lenient `.get()`), and the real captured
+data is best-effort, not a formal spec pulled from documentation. Treat
+this the same way as `_redact.py`: best-effort coverage of the tools this
+environment is actually observed calling, not a guaranteed-complete or
+guaranteed-accurate schema for any of them.
+
+**Known, explicit gap**: this module only covers tools actually observed
+in this environment's corpus. A tool_call to anything outside
+`TOOL_SCHEMAS` means the training example gets rendered with `tools=None`
+entirely (see `to_text()`'s uncovered-tool guard) rather than a `tools=`
+list that's missing what the assistant then calls -- omitting a tool from
+the shown schema while still training the model to call it is the exact
+anti-pattern schema-conditioning exists to prevent.
+"""
+
+from __future__ import annotations
+
+
+def _fn(name: str, description: str, properties: dict, required: list) -> dict:
+    return {
+        "type": "function",
+        "function": {
+            "name": name,
+            "description": description,
+            "parameters": {"type": "object", "properties": properties, "required": required},
+        },
+    }
+
+
+_STR = {"type": "string"}
+_INT = {"type": "integer"}
+_BOOL = {"type": "boolean"}
+
+
+def _arr(item_type) -> dict:
+    return {"type": "array", "items": item_type}
+
+
+TOOL_SCHEMAS: dict = {
+    # -- 11 vanilla Claude Code tools (fields from _activity.py's own
+    # _summarize_*/_full_* extractors) --
+    "Read": _fn(
+        "Read", "Read a file.",
+        {"file_path": _STR, "offset": _INT, "limit": _INT},
+        ["file_path"],
+    ),
+    "Edit": _fn(
+        "Edit", "Edit a file via exact string replacement.",
+        {"file_path": _STR, "old_string": _STR, "new_string": _STR, "replace_all": _BOOL},
+        ["file_path", "old_string", "new_string"],
+    ),
+    "Write": _fn(
+        "Write", "Write a file, creating or overwriting it.",
+        {"file_path": _STR, "content": _STR},
+        ["file_path", "content"],
+    ),
+    "Bash": _fn(
+        "Bash", "Run a shell command.",
+        {"command": _STR, "timeout": _INT, "description": _STR, "run_in_background": _BOOL},
+        ["command"],
+    ),
+    "Grep": _fn(
+        "Grep", "Search file contents (ripgrep-backed).",
+        {
+            "pattern": _STR, "path": _STR, "-i": _BOOL, "-n": _BOOL,
+            "-A": _INT, "-B": _INT, "-C": _INT, "output_mode": _STR,
+            "type": _STR, "head_limit": _INT, "multiline": _BOOL,
+        },
+        ["pattern"],
+    ),
+    "Glob": _fn(
+        "Glob", "Find files by glob pattern.",
+        {"pattern": _STR, "path": _STR},
+        ["pattern"],
+    ),
+    "Agent": _fn(
+        "Agent", "Dispatch a sub-agent task.",
+        {"description": _STR, "prompt": _STR, "subagent_type": _STR},
+        ["description", "prompt", "subagent_type"],
+    ),
+    "TodoWrite": _fn(
+        "TodoWrite", "Replace the full todo list.",
+        {"todos": _arr({"type": "object", "properties": {
+            "content": _STR, "status": _STR, "activeForm": _STR,
+        }})},
+        ["todos"],
+    ),
+    "WebSearch": _fn(
+        "WebSearch", "Search the web.",
+        {"query": _STR, "allowed_domains": _arr(_STR), "blocked_domains": _arr(_STR)},
+        ["query"],
+    ),
+    "WebFetch": _fn(
+        "WebFetch", "Fetch a URL and extract content per a prompt.",
+        {"url": _STR, "prompt": _STR},
+        ["url", "prompt"],
+    ),
+    "AskUserQuestion": _fn(
+        "AskUserQuestion", "Ask the user a structured multiple-choice question.",
+        {"questions": _arr({"type": "object", "properties": {
+            "question": _STR, "header": _STR, "multiSelect": _BOOL,
+            "options": _arr({"type": "object", "properties": {
+                "label": _STR, "description": _STR,
+            }}),
+        }})},
+        ["questions"],
+    ),
+    # -- 13 harness-native tools, sampled from real tool_input values in
+    # this environment's own corpus (data/activity.db) --
+    "DesignSync": _fn(
+        "DesignSync", "Sync a design doc.",
+        {"method": _STR, "projectId": _STR},
+        ["method", "projectId"],
+    ),
+    "EnterPlanMode": _fn("EnterPlanMode", "Enter plan mode.", {}, []),
+    "ExitPlanMode": _fn(
+        "ExitPlanMode", "Exit plan mode with a plan for approval.",
+        {"plan": _STR},
+        ["plan"],
+    ),
+    "Monitor": _fn(
+        "Monitor", "Monitor a long-running command, streaming matching lines as events.",
+        {"description": _STR, "command": _STR, "timeout_ms": _INT, "persistent": _BOOL},
+        ["description", "command"],
+    ),
+    "ScheduleWakeup": _fn(
+        "ScheduleWakeup", "Schedule a future wakeup.",
+        {"delaySeconds": _INT, "reason": _STR, "prompt": _STR},
+        ["delaySeconds", "reason", "prompt"],
+    ),
+    "SendMessage": _fn(
+        "SendMessage", "Send a message to another agent.",
+        {"to": _STR, "message": _STR, "summary": _STR},
+        ["to", "message"],
+    ),
+    "Skill": _fn(
+        "Skill", "Invoke a named skill.",
+        {"skill": _STR, "args": _STR},
+        ["skill"],
+    ),
+    "TaskCreate": _fn(
+        "TaskCreate", "Create a tracked background task.",
+        {"subject": _STR, "description": _STR},
+        ["subject"],
+    ),
+    "TaskList": _fn("TaskList", "List tracked tasks.", {}, []),
+    "TaskOutput": _fn(
+        "TaskOutput", "Get a tracked task's output.",
+        {"task_id": _STR, "block": _BOOL, "timeout": _INT},
+        ["task_id"],
+    ),
+    "TaskStop": _fn(
+        "TaskStop", "Stop a running tracked task.",
+        {"task_id": _STR},
+        ["task_id"],
+    ),
+    "TaskUpdate": _fn(
+        "TaskUpdate", "Update a tracked task's status.",
+        # NB: `taskId`, not `task_id` -- a real inconsistency observed in
+        # the corpus versus TaskOutput/TaskStop, kept faithfully rather
+        # than silently normalized.
+        {"taskId": _STR, "status": _STR},
+        ["taskId", "status"],
+    ),
+    "ToolSearch": _fn(
+        "ToolSearch", "Search for a deferred tool by name or keyword.",
+        {"query": _STR, "max_results": _INT},
+        ["query"],
+    ),
+}
+
+ALL_TOOLS: list = sorted(TOOL_SCHEMAS.values(), key=lambda t: t["function"]["name"])

blitz_cli-0.6.0/blitz_cli/templates/codeassist_Dockerfile.train.tmpl

@@ -0,0 +1,26 @@
+# Heavy QLoRA trainer for the code-assist corpus (data/activity.db). Kept as
+# a SEPARATE image from the lightweight watcher (Dockerfile) -- the always-on
+# watcher just polls the filesystem and doesn't need a GPU or the ML stack;
+# baking that stack into the same image it runs continuously would bloat a
+# daemon meant to idle for weeks. This image is only built/run on demand for
+# `blitz code-assist train`.
+#
+# CUDA 13.0, not 12.x: needed for Blackwell (sm_120, e.g. RTX 50-series)
+# kernels -- torch/bitsandbytes built against older CUDA raise "no kernel
+# image is available for execution on the device" on those cards. Verified
+# against a real RTX 5070 Ti (driver reports CUDA 13.0 as its max supported
+# version, so this matches rather than exceeds the host).
+FROM pytorch/pytorch:2.13.0-cuda13.0-cudnn9-runtime
+
+WORKDIR /workspace
+ENV PYTHONUNBUFFERED=1 HF_HOME=/workspace/.hf
+
+COPY requirements-train.txt .
+# --break-system-packages: this base image ships a PEP 668 "externally
+# managed" system Python (unlike the old conda-based 2.5.1 image), and there's
+# no venv to install into here -- it's a single-purpose training container.
+RUN pip install --no-cache-dir --break-system-packages -r requirements-train.txt
+
+COPY train.py .
+
+CMD ["python", "train.py"]

{blitz_cli-0.5.0 → blitz_cli-0.6.0}/blitz_cli/templates/codeassist_Makefile.tmpl

@@ -43,6 +43,7 @@ RUN = docker run -d \
 TRAIN_RUN = docker run --rm --gpus all \
 	-v $(CURDIR):/workspace \
 	-e BASE_MODEL_HF=$(BASE_MODEL_HF) \
+	-e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
 	$(TRAIN_IMAGE)
 
 .PHONY: build run logs stop shell build-train train

{blitz_cli-0.5.0 → blitz_cli-0.6.0}/blitz_cli/templates/codeassist_README.md.tmpl

@@ -79,7 +79,10 @@ are already chat-formatted, tool-calling trajectories. Override the base
 model per-run with `BASE_MODEL_HF=<hf-id>`, or bake in a different default at
 scaffold time with `blitz scaffold code-assist --base-model <hf-id>`. Other
 overrides: `SEQ_LEN`, `PARAMS_B` (drives LoRA rank/batch-size defaults),
-`EVAL_FRACTION`.
+`EVAL_FRACTION`, and `PER_DEVICE_BATCH_SIZE` (overrides just the batch size
+from `PARAMS_B`'s bucket, independent of LoRA rank -- rescales `grad_accum`
+to hold the effective batch size steady, so it's a free memory/step-count
+tradeoff with no effect on training dynamics, unlike `SEQ_LEN`).
 
 Each run is tracked in `training.db` (MLflow's SQLite backend) and saves its
 adapter to `adapter-v<N>` (N = that run's registered model version), so

blitz_cli-0.6.0/blitz_cli/templates/codeassist_requirements_train.txt.tmpl

@@ -0,0 +1,17 @@
+# QLoRA SFT stack for training on the code-assist corpus (data/activity.db).
+# torch is intentionally omitted -- the Dockerfile.train's pytorch/cuda base
+# image ships the matching CUDA build, and reinstalling risks pulling a
+# CPU-only wheel or one built for the wrong CUDA/arch.
+transformers==4.48.0
+trl==0.13.0
+peft==0.14.0
+accelerate==1.2.1
+# 0.45.0 has no compiled kernels for Blackwell (sm_120, e.g. RTX 50-series) --
+# 4-bit quantization raises "no kernel image is available for execution on
+# the device" on those cards. 0.49.2 verified working on a real RTX 5070 Ti.
+bitsandbytes==0.49.2
+datasets==3.2.0
+sentencepiece==0.2.0
+# Local SQLite tracking/registry for this run (sqlite:///training.db) -- no
+# `mlflow server` process, just the tracking client + SQLAlchemy backend.
+mlflow==2.19.0

{blitz_cli-0.5.0 → blitz_cli-0.6.0}/blitz_cli/templates/codeassist_train.py.tmpl

@@ -6,10 +6,36 @@ Defaults to Qwen2.5-Coder-14B-Instruct: our training examples are already
 chat-formatted, tool-calling trajectories (native Claude Code tool names --
 Read/Edit/Write/Bash/...), and an instruction-tuned checkpoint already
 understands a chat template and basic tool-calling behavior, which is the
-right starting point for further SFT rather than a raw base model. Override
-via the BASE_MODEL_HF env var (also settable at scaffold time with
+right starting point for further SFT rather than a raw base model. 14B
+assumes real headroom (a rented 80GB-class GPU) -- on a 16GB-class consumer
+card, 7B was the safe default instead (14B OOM'd during LoRA adapter setup,
+before any forward/backward pass, on a real RTX 5070 Ti). Override the base
+model via the BASE_MODEL_HF env var (also settable at scaffold time with
 `blitz scaffold code-assist --base-model <hf-id>`, which bakes a new default
-into this project's Makefile).
+into this project's Makefile) -- and PARAMS_B to match, since it drives the
+LoRA rank/batch-size defaults in _scaffold.build_config.
+
+SEQ_LEN defaults to 16384, not 2048, once real headroom is available --
+measured (not estimated) against the real corpus + the real Qwen tokenizer
+with the full tool-schema block from _tool_schemas.py included on every
+example (see below): at seq_len=2048 the schema tax alone pushes 77.6% of
+turns over budget (vs. 42.5% with no schema at all -- a +35.1pp hit); at
+16384 that gap shrinks to +0.8pp (7.3% vs. 6.5%), i.e. the schema is nearly
+free relative to its own baseline at this length. On a memory-constrained
+card, packing still splits a turn longer than SEQ_LEN across separate
+training steps with no attention connection between the pieces -- raise or
+lower this based on actual available VRAM, using the measured table above
+as a guide rather than guessing.
+
+If you're still tight on memory after that: PER_DEVICE_BATCH_SIZE overrides
+build_config's params_b-driven batch size directly, independent of LoRA
+rank (unlike PARAMS_B, which moves both together as one bucket). Lowering
+it rescales grad_accum to hold the effective batch size steady, so --
+unlike SEQ_LEN -- this one is free: no training-dynamics or data-quality
+tradeoff, just more (smaller) steps for the same total memory footprint per
+step. On real hardware with unmeasured headroom, start with a conservative
+guess and watch `nvidia-smi` on the first few steps rather than trusting a
+guessed number.
 
 Also tracks the run in ./training.db (MLflow's SQLite backend -- no `mlflow
 server` process, just a local file): per-step loss history and a registered
@@ -33,15 +59,17 @@ from peft import LoraConfig, prepare_model_for_kbit_training
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from trl import SFTConfig, SFTTrainer
 
-from blitz_cli import _dataset, _scaffold
+from blitz_cli import _dataset, _scaffold, _tool_schemas
 
 mlflow.set_tracking_uri("sqlite:///training.db")
 
 _BASE_MODEL_HF = os.environ.get("BASE_MODEL_HF", "Qwen/Qwen2.5-Coder-14B-Instruct")
 _PARAMS_B = float(os.environ.get("PARAMS_B", "14"))
-_SEQ_LEN = _scaffold.clamp_seq_len(int(os.environ.get("SEQ_LEN", "4096")))
+_SEQ_LEN = _scaffold.clamp_seq_len(int(os.environ.get("SEQ_LEN", "16384")))
 _EVAL_FRACTION = float(os.environ.get("EVAL_FRACTION", "0.1"))
 _DB_PATH = os.environ.get("ACTIVITY_DB_PATH", "data/activity.db")
+_SAVE_STEPS = int(os.environ.get("SAVE_STEPS", "50"))
+_ALL_TOOL_NAMES = frozenset(_tool_schemas.TOOL_SCHEMAS)
 
 cfg = _scaffold.build_config(
     workflow="code-assist",
@@ -50,7 +78,15 @@ cfg = _scaffold.build_config(
     seq_len=_SEQ_LEN,
     eval_fraction=_EVAL_FRACTION,
 )
-print(f"Base model: {cfg['base_model_hf']}  seq_len={cfg['seq_len']}  params_b={cfg['params_b']}")
+
+_batch_override = os.environ.get("PER_DEVICE_BATCH_SIZE")
+if _batch_override:
+    cfg = _scaffold.apply_batch_override(cfg, int(_batch_override))
+
+print(
+    f"Base model: {cfg['base_model_hf']}  seq_len={cfg['seq_len']}  params_b={cfg['params_b']}  "
+    f"per_device_batch_size={cfg['per_device_batch_size']}  grad_accum={cfg['grad_accum']}"
+)
 
 if not Path(_DB_PATH).exists():
     raise SystemExit(
@@ -123,10 +159,36 @@ def to_text(example):
     # the model's own chat template. Assistant messages carry `tool_calls`;
     # tool-role messages carry `tool_call_id`. If tool-call rendering doesn't
     # look right for the exact Qwen revision in use, check its chat template
-    # in tokenizer_config.json -- we don't pass a `tools=` schema list here
-    # since Claude Code's own transcripts don't record each tool's JSON
-    # input_schema, only the calls actually made.
-    return {"text": tokenizer.apply_chat_template(example["messages"], tokenize=False)}
+    # in tokenizer_config.json.
+    #
+    # `arguments` is stored as a JSON *string* (the real OpenAI wire format --
+    # function.arguments is a string there too), but Qwen's chat template
+    # does `tool_call.arguments | tojson`, which assumes a raw object and
+    # JSON-encodes it itself. Feeding it an already-stringified value double-
+    # encodes it (confirmed via a real corpus example: `"arguments":
+    # "{\"file_path\": ...}"` instead of a clean nested object), producing
+    # garbled training text and abnormally high loss. Parse it back to a
+    # dict here, at render time only -- the stored messages_json keeps the
+    # portable OpenAI-string convention for any other consumer.
+    #
+    # tools=: pass the full fixed tool-schema list (_tool_schemas.ALL_TOOLS)
+    # on every example, so the model learns to read a given schema and
+    # comply with it -- not just memorize this environment's fixed tool
+    # vocabulary. If an example calls a tool outside _tool_schemas.TOOL_SCHEMAS,
+    # fall back to tools=None entirely rather than show a list that's
+    # missing what the assistant then calls -- that mismatch is worse than
+    # showing no schema at all.
+    messages = json.loads(json.dumps(example["messages"]))  # deep copy, don't mutate the dataset row
+    tool_names = set()
+    for m in messages:
+        for tc in m.get("tool_calls") or []:
+            fn = tc.get("function") or {}
+            if fn.get("name"):
+                tool_names.add(fn["name"])
+            if isinstance(fn.get("arguments"), str):
+                fn["arguments"] = json.loads(fn["arguments"])
+    tools = _tool_schemas.ALL_TOOLS if tool_names <= _ALL_TOOL_NAMES else None
+    return {"text": tokenizer.apply_chat_template(messages, tools=tools, tokenize=False)}
 
 
 dataset = dataset.map(to_text, remove_columns=dataset.column_names)
@@ -141,7 +203,13 @@ sft_config = SFTConfig(
     packing=True,
     bf16=True,
     logging_steps=10,
-    save_strategy="epoch",
+    # steps, not epoch: LoRA checkpoints are small and cheap to write, so
+    # there's no real cost to checkpointing often -- only upside if a run
+    # gets interrupted (a spot instance reclaim, a crash, anything) partway
+    # through what could be an hours-long epoch on a real corpus. Tune
+    # SAVE_STEPS once real steps/sec is known for a ~5-10 minute cadence.
+    save_strategy="steps",
+    save_steps=_SAVE_STEPS,
     lr_scheduler_type="cosine",
     warmup_ratio=0.03,
     gradient_checkpointing=True,

{blitz_cli-0.5.0 → blitz_cli-0.6.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "blitz-cli"
-version = "0.5.0"
+version = "0.6.0"
 description = "Developer CLI for Blitz: local-first tracing/init, codebase instrumentation scan, and scaffolds for both a QLoRA workflow trainer and a coding-agent session-to-fine-tuning-corpus watcher."
 readme = "README.md"
 license = { text = "MIT" }

{blitz_cli-0.5.0 → blitz_cli-0.6.0}/tests/test_scaffold.py

@@ -36,7 +36,7 @@ def test_build_config_carries_base_model_and_workflow():
 
 def test_seq_len_is_clamped():
     assert _cfg(seq_len=4096)["seq_len"] == 4096
-    assert _cfg(seq_len=999999)["seq_len"] == 8192  # upper clamp
+    assert _cfg(seq_len=999999)["seq_len"] == 16384  # upper clamp
     assert _cfg(seq_len=10)["seq_len"] == 1024       # lower clamp
     assert _cfg(seq_len=0)["seq_len"] == 2048         # falsy -> default
 
@@ -45,6 +45,8 @@ def test_clamp_seq_len_helper():
     assert _scaffold.clamp_seq_len(1200) == 1200
     assert _scaffold.clamp_seq_len(None) == 2048
     assert _scaffold.clamp_seq_len(50) == 1024
+    assert _scaffold.clamp_seq_len(16384) == 16384
+    assert _scaffold.clamp_seq_len(999999) == 16384
 
 
 def test_gated_license_warning():
@@ -64,6 +66,31 @@ def test_build_config_eval_fraction_override():
     assert cfg["eval_fraction"] == 0.05
 
 
+def test_apply_batch_override_holds_effective_batch_size_steady():
+    # params_b=7 bucket: per_device_batch_size=2, grad_accum=8 -> effective 16
+    cfg = _cfg(params_b=7.0)
+    overridden = _scaffold.apply_batch_override(cfg, 1)
+    assert overridden["per_device_batch_size"] == 1
+    assert overridden["grad_accum"] == 16  # 2*8 / 1, effective batch unchanged
+    # only batch/accum change -- lora_r and everything else carry over as-is
+    assert overridden["lora_r"] == cfg["lora_r"]
+    assert overridden["base_model_hf"] == cfg["base_model_hf"]
+
+
+def test_apply_batch_override_rounds_when_not_evenly_divisible():
+    cfg = _cfg(params_b=3.0)  # per_device_batch_size=4, grad_accum=4 -> effective 16
+    overridden = _scaffold.apply_batch_override(cfg, 3)
+    assert overridden["per_device_batch_size"] == 3
+    assert overridden["grad_accum"] == 5  # round(16/3) = round(5.33) = 5
+
+
+def test_apply_batch_override_never_drops_grad_accum_below_one():
+    cfg = _cfg(params_b=3.0)  # effective batch 16
+    overridden = _scaffold.apply_batch_override(cfg, 64)  # way bigger than effective batch
+    assert overridden["grad_accum"] == 1
+    assert overridden["per_device_batch_size"] == 64
+
+
 def test_render_writes_a_runnable_skeleton_without_config(tmp_path):
     out = tmp_path / "train"
     _scaffold.render(out, "wf")
@@ -185,6 +212,55 @@ def test_render_code_assist_train_py_reads_activity_db_and_defaults_to_qwen(tmp_
     assert "apply_chat_template" in train_py
 
 
+def test_render_code_assist_train_py_seq_len_defaults_to_16384(tmp_path):
+    out = tmp_path / "code-assist"
+    _scaffold.render_code_assist(out)
+    train_py = (out / "train.py").read_text()
+    assert 'os.environ.get("SEQ_LEN", "16384")' in train_py
+
+
+def test_render_code_assist_train_py_wires_tools_into_chat_template(tmp_path):
+    out = tmp_path / "code-assist"
+    _scaffold.render_code_assist(out)
+    train_py = (out / "train.py").read_text()
+    assert "_tool_schemas" in train_py
+    assert "tool_names <= _ALL_TOOL_NAMES" in train_py
+    assert "tools=tools" in train_py
+
+
+def test_render_code_assist_train_py_checkpoints_by_steps_not_epoch(tmp_path):
+    out = tmp_path / "code-assist"
+    _scaffold.render_code_assist(out)
+    train_py = (out / "train.py").read_text()
+    assert 'save_strategy="steps"' in train_py
+    assert "save_steps=_SAVE_STEPS" in train_py
+    assert 'save_strategy="epoch"' not in train_py
+
+
+def test_render_code_assist_train_py_parses_tool_call_arguments_before_templating(tmp_path):
+    # regression test: _codeassist.turn_to_messages stores tool_calls[i].function
+    # .arguments as a JSON *string* (correct -- matches the real OpenAI wire
+    # format), but Qwen's chat template does `arguments | tojson`, which
+    # assumes a raw object and double-encodes an already-stringified value.
+    # Confirmed on a real corpus example: training text ended up with escaped
+    # JSON-inside-a-string garbage (e.g. `"arguments": "{\"file_path\": ...}"`
+    # instead of a clean nested object), which correlated with abnormally high
+    # training loss. to_text must parse arguments back to a dict first.
+    out = tmp_path / "code-assist"
+    _scaffold.render_code_assist(out)
+    train_py = (out / "train.py").read_text()
+    assert 'isinstance(fn.get("arguments"), str)' in train_py
+    assert 'fn["arguments"] = json.loads(fn["arguments"])' in train_py
+
+
+def test_render_code_assist_train_py_wires_up_batch_override(tmp_path):
+    out = tmp_path / "code-assist"
+    _scaffold.render_code_assist(out)
+    train_py = (out / "train.py").read_text()
+    assert 'os.environ.get("PER_DEVICE_BATCH_SIZE")' in train_py
+    assert "_scaffold.apply_batch_override(cfg, int(_batch_override))" in train_py
+
+
 def test_render_code_assist_makefile_substitutes_interval_base_model_and_bind_mounts(tmp_path):
     out = tmp_path / "code-assist"
     _scaffold.render_code_assist(out, interval=30, base_model="Qwen/Qwen2.5-Coder-7B-Instruct")
@@ -198,15 +274,22 @@ def test_render_code_assist_makefile_substitutes_interval_base_model_and_bind_mo
     assert "--restart unless-stopped" in mk
     assert "--gpus all" in mk
     assert "Dockerfile.train" in mk
+    # free fragmentation-avoidance win for the trainer -- verified on real
+    # hardware to help (though not by itself sufficient for a 14B model on a
+    # 16GB card, see the 7B default above)
+    assert "PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True" in mk
     for target in ("build:", "run:", "logs:", "stop:", "shell:", "build-train:", "train:"):
         assert target in mk
 
 
-def test_render_code_assist_defaults_to_qwen_14b_instruct(tmp_path):
+def test_render_code_assist_defaults_to_qwen_7b_instruct(tmp_path):
+    # 7B, not 14B: the size class that comfortably fits a QLoRA run on a
+    # 16GB-class consumer card -- 14B was found (on real hardware) to OOM
+    # during LoRA adapter setup, before any training step even ran.
     out = tmp_path / "code-assist"
     _scaffold.render_code_assist(out)
     mk = (out / "Makefile").read_text()
-    assert "BASE_MODEL_HF ?= Qwen/Qwen2.5-Coder-14B-Instruct" in mk
+    assert "BASE_MODEL_HF ?= Qwen/Qwen2.5-Coder-7B-Instruct" in mk
 
 
 def test_render_code_assist_readme_substitutes_interval_and_base_model(tmp_path):
@@ -229,7 +312,7 @@ def test_render_code_assist_appends_blitz_cli_pin_to_both_requirements_files(tmp
     assert "pip install --no-cache-dir -r requirements.txt" in dockerfile
     assert "_vendor" not in dockerfile
     train_dockerfile = (out / "Dockerfile.train").read_text()
-    assert "pip install --no-cache-dir -r requirements-train.txt" in train_dockerfile
+    assert "pip install --no-cache-dir --break-system-packages -r requirements-train.txt" in train_dockerfile
     assert "_vendor" not in train_dockerfile
 
 

blitz_cli-0.6.0/tests/test_tool_schemas.py

@@ -0,0 +1,117 @@
+import json
+
+from blitz_cli import _tool_schemas
+
+# The 24 tools this module is expected to cover: the 11 vanilla Claude Code
+# tools plus the 13 harness-native tools sampled from the real corpus.
+_EXPECTED_TOOLS = {
+    "Read", "Edit", "Write", "Bash", "Grep", "Glob", "Agent", "TodoWrite",
+    "WebSearch", "WebFetch", "AskUserQuestion",
+    "DesignSync", "EnterPlanMode", "ExitPlanMode", "Monitor", "ScheduleWakeup",
+    "SendMessage", "Skill", "TaskCreate", "TaskList", "TaskOutput", "TaskStop",
+    "TaskUpdate", "ToolSearch",
+}
+
+
+def test_tool_schemas_covers_all_24_expected_tools():
+    assert set(_tool_schemas.TOOL_SCHEMAS) == _EXPECTED_TOOLS
+
+
+def test_every_schema_has_the_expected_shape_and_round_trips_through_json():
+    for name, schema in _tool_schemas.TOOL_SCHEMAS.items():
+        assert schema["type"] == "function"
+        fn = schema["function"]
+        assert fn["name"] == name
+        assert isinstance(fn["description"], str) and fn["description"]
+        params = fn["parameters"]
+        assert params["type"] == "object"
+        assert isinstance(params["properties"], dict)
+        assert isinstance(params["required"], list)
+        # every required field must actually be a declared property
+        assert set(params["required"]) <= set(params["properties"])
+        # plain JSON-able data, no stray non-serializable values
+        assert json.loads(json.dumps(schema)) == schema
+
+
+def test_read_matches_activity_py_capture():
+    props = _tool_schemas.TOOL_SCHEMAS["Read"]["function"]["parameters"]
+    assert set(props["properties"]) == {"file_path", "offset", "limit"}
+    assert props["required"] == ["file_path"]
+
+
+def test_edit_matches_activity_py_capture():
+    props = _tool_schemas.TOOL_SCHEMAS["Edit"]["function"]["parameters"]
+    assert set(props["properties"]) == {"file_path", "old_string", "new_string", "replace_all"}
+    assert set(props["required"]) == {"file_path", "old_string", "new_string"}
+
+
+def test_write_matches_activity_py_capture():
+    props = _tool_schemas.TOOL_SCHEMAS["Write"]["function"]["parameters"]
+    assert set(props["properties"]) == {"file_path", "content"}
+    assert set(props["required"]) == {"file_path", "content"}
+
+
+def test_bash_matches_activity_py_capture():
+    props = _tool_schemas.TOOL_SCHEMAS["Bash"]["function"]["parameters"]
+    assert set(props["properties"]) == {"command", "timeout", "description", "run_in_background"}
+    assert props["required"] == ["command"]
+
+
+def test_grep_matches_activity_py_capture():
+    props = _tool_schemas.TOOL_SCHEMAS["Grep"]["function"]["parameters"]
+    assert set(props["properties"]) == {
+        "pattern", "path", "-i", "-n", "-A", "-B", "-C",
+        "output_mode", "type", "head_limit", "multiline",
+    }
+    assert props["required"] == ["pattern"]
+
+
+def test_agent_requires_all_three_fields():
+    props = _tool_schemas.TOOL_SCHEMAS["Agent"]["function"]["parameters"]
+    assert set(props["properties"]) == {"description", "prompt", "subagent_type"}
+    assert set(props["required"]) == {"description", "prompt", "subagent_type"}
+
+
+def test_todowrite_is_a_full_list_replace_array():
+    props = _tool_schemas.TOOL_SCHEMAS["TodoWrite"]["function"]["parameters"]
+    assert set(props["properties"]) == {"todos"}
+    todos_schema = props["properties"]["todos"]
+    assert todos_schema["type"] == "array"
+    assert set(todos_schema["items"]["properties"]) == {"content", "status", "activeForm"}
+
+
+def test_askuserquestion_nested_options_shape():
+    props = _tool_schemas.TOOL_SCHEMAS["AskUserQuestion"]["function"]["parameters"]
+    questions_schema = props["properties"]["questions"]
+    item_props = questions_schema["items"]["properties"]
+    assert set(item_props) == {"question", "header", "multiSelect", "options"}
+    option_props = item_props["options"]["items"]["properties"]
+    assert set(option_props) == {"label", "description"}
+
+
+def test_taskupdate_uses_taskid_camelcase_not_task_id():
+    # a real inconsistency observed in the corpus vs TaskOutput/TaskStop's
+    # task_id -- kept faithfully rather than silently normalized.
+    props = _tool_schemas.TOOL_SCHEMAS["TaskUpdate"]["function"]["parameters"]
+    assert "taskId" in props["properties"]
+    assert "task_id" not in props["properties"]
+
+
+def test_taskoutput_and_taskstop_use_task_id_snakecase():
+    for name in ("TaskOutput", "TaskStop"):
+        props = _tool_schemas.TOOL_SCHEMAS[name]["function"]["parameters"]
+        assert "task_id" in props["properties"]
+
+
+def test_zero_arg_tools_have_empty_properties_and_required():
+    for name in ("EnterPlanMode", "TaskList"):
+        props = _tool_schemas.TOOL_SCHEMAS[name]["function"]["parameters"]
+        assert props["properties"] == {}
+        assert props["required"] == []
+
+
+def test_all_tools_is_sorted_by_name_with_one_entry_per_schema():
+    names = [t["function"]["name"] for t in _tool_schemas.ALL_TOOLS]
+    assert names == sorted(names)
+    assert len(_tool_schemas.ALL_TOOLS) == len(_tool_schemas.TOOL_SCHEMAS)
+    assert set(names) == _EXPECTED_TOOLS

blitz_cli-0.5.0/blitz_cli/templates/codeassist_Dockerfile.train.tmpl

@@ -1,17 +0,0 @@
-# Heavy QLoRA trainer for the code-assist corpus (data/activity.db). Kept as
-# a SEPARATE image from the lightweight watcher (Dockerfile) -- the always-on
-# watcher just polls the filesystem and doesn't need a GPU or the ML stack;
-# baking that stack into the same image it runs continuously would bloat a
-# daemon meant to idle for weeks. This image is only built/run on demand for
-# `blitz code-assist train`.
-FROM pytorch/pytorch:2.5.1-cuda12.1-cudnn9-runtime
-
-WORKDIR /workspace
-ENV PYTHONUNBUFFERED=1 HF_HOME=/workspace/.hf
-
-COPY requirements-train.txt .
-RUN pip install --no-cache-dir -r requirements-train.txt
-
-COPY train.py .
-
-CMD ["python", "train.py"]

blitz_cli-0.5.0/blitz_cli/templates/codeassist_requirements_train.txt.tmpl

@@ -1,15 +0,0 @@
-# QLoRA SFT stack for training on the code-assist corpus (data/activity.db),
-# pinned to the same mutually-compatible set as the QLoRA workflow trainer's
-# requirements.txt (CUDA 12.1, early 2026). torch is intentionally omitted --
-# the Dockerfile.train's pytorch/cuda base image ships the matching CUDA
-# build, and reinstalling risks pulling a CPU-only wheel.
-transformers==4.48.0
-trl==0.13.0
-peft==0.14.0
-accelerate==1.2.1
-bitsandbytes==0.45.0
-datasets==3.2.0
-sentencepiece==0.2.0
-# Local SQLite tracking/registry for this run (sqlite:///training.db) -- no
-# `mlflow server` process, just the tracking client + SQLAlchemy backend.
-mlflow==2.19.0