8be8d003e7
Script + pytest tests to verify MLX inference works at CorridorKey's training resolution. Uses samples/ by default, synthetic fallback. Reports timing, peak memory, output diagnostics. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
118 lines
4.7 KiB
Markdown
118 lines
4.7 KiB
Markdown
---
|
||
title: "feat: Add 2048 smoke test"
|
||
type: feat
|
||
date: 2026-03-01
|
||
---
|
||
|
||
# Add 2048 Smoke Test
|
||
|
||
## Overview
|
||
|
||
Validate MLX model runs end-to-end at CorridorKey's native 2048×2048 resolution on Apple Silicon. Smoke/stability check — not a new parity campaign.
|
||
|
||
## Problem Statement
|
||
|
||
All existing tests/scripts default to 256–512. No automated path confirms 2048 actually works. The model was *trained* at 2048, so pos_embed interpolation is a no-op at that resolution, but we haven't exercised the full pipeline there.
|
||
|
||
## Proposed Solution
|
||
|
||
One new script + one skipable pytest test. Reuse `CorridorKeyMLXEngine` (already defaults to `img_size=2048`). Minimal additions.
|
||
|
||
## Deliverables
|
||
|
||
### 1. `scripts/smoke_2048.py`
|
||
|
||
Mirrors `scripts/smoke_engine.py` pattern but targeted at 2048 with richer diagnostics.
|
||
|
||
**CLI args:**
|
||
- `--checkpoint PATH` (default: `checkpoints/corridorkey_mlx.safetensors`)
|
||
- `--img-size INT` (default: **2048** — this script's whole point)
|
||
- `--image PATH` (optional — uses synthetic input if omitted)
|
||
- `--hint PATH` (optional — uses synthetic hint if omitted)
|
||
- `--output-dir DIR` (default: `output/smoke_2048/`)
|
||
- `--save-outputs / --no-save-outputs` (default: save)
|
||
- `--seed INT` (default: 42)
|
||
- `--compile / --no-compile` (default: True — engine default)
|
||
|
||
**Behavior:**
|
||
1. Generate or load 2048×2048 RGB + hint inputs
|
||
2. Instantiate `CorridorKeyMLXEngine(checkpoint_path, img_size, compile)`
|
||
3. Reset peak memory via `mx.metal.reset_peak_memory()` (if available, try/except)
|
||
4. Time `engine.process_frame(rgb, mask)` with wall-clock
|
||
5. Read peak memory via `mx.metal.get_peak_memory()` (if available)
|
||
6. Report: input shape, output shapes, dtypes, elapsed time, peak memory (MB)
|
||
7. Diagnostic: min/max/mean of alpha+fg, NaN/Inf check, flag all-zeros/all-ones
|
||
8. Optionally save alpha.png, fg.png, comp.png
|
||
9. Catch `RuntimeError`/`MemoryError` with actionable error message
|
||
|
||
**Synthetic input generation:**
|
||
```python
|
||
rng = np.random.default_rng(seed)
|
||
rgb = rng.integers(0, 256, (2048, 2048, 3), dtype=np.uint8)
|
||
# Circular gradient hint — exercises the path better than random noise
|
||
mask = <simple radial gradient uint8>
|
||
```
|
||
|
||
### 2. `tests/test_smoke_2048.py`
|
||
|
||
**Two tests:**
|
||
|
||
1. `test_smoke_2048_full` — loads checkpoint, runs engine at 2048, asserts shapes + no NaN. Marked `@pytest.mark.skipif(not HAS_CHECKPOINT)` AND `@pytest.mark.slow`. Won't run in normal `uv run pytest`; run via `uv run pytest -m slow`.
|
||
|
||
2. `test_smoke_2048_wiring` — lightweight, no checkpoint. Verifies `GreenFormer(img_size=2048)` constructs OK and produces correct output shapes with random weights at 2048 (or smaller like 256 if too heavy without checkpoint). Runs in normal suite.
|
||
|
||
### 3. README update
|
||
|
||
Add "2048 Smoke Test" section after the existing "Smoke test" subsection:
|
||
|
||
```markdown
|
||
### 2048 smoke test
|
||
|
||
Validates full end-to-end inference at CorridorKey's native 2048 resolution.
|
||
Uses synthetic inputs by default — no test images required.
|
||
|
||
```bash
|
||
uv run python scripts/smoke_2048.py
|
||
```
|
||
|
||
To use real images:
|
||
```bash
|
||
uv run python scripts/smoke_2048.py --image shot.png --hint hint.png
|
||
```
|
||
|
||
Reports timing, peak memory, output shapes, and value-range diagnostics.
|
||
This is an execution check, not a 2048 parity validation.
|
||
```
|
||
|
||
## Technical Considerations
|
||
|
||
- **Memory**: 2048×2048×4 float32 ≈ 64MB input, but backbone + decoder intermediates will be much larger. Engine already handles this; we just report peak.
|
||
- **Compile**: Engine defaults `compile=True`. First call at 2048 will be slow (compile cost). The smoke script is a single-run tool so the compile overhead is included in timing.
|
||
- **Pos_embed**: At 2048, no interpolation needed (trained at 2048). This is the easiest resolution to test.
|
||
- **`mx.metal`**: May not be available in all environments. Wrap in try/except.
|
||
|
||
## Acceptance Criteria
|
||
|
||
- [x] `scripts/smoke_2048.py` runs successfully with checkpoint
|
||
- [x] Reports: input size, output shapes, elapsed time, peak memory
|
||
- [x] Detects NaN/Inf and flags suspicious outputs
|
||
- [x] Supports both synthetic and user-supplied inputs
|
||
- [x] `test_smoke_2048_full` passes when checkpoint present + `-m slow`
|
||
- [x] `test_smoke_2048_wiring` passes in normal test suite
|
||
- [x] README documents the smoke test
|
||
- [x] Existing 512 tests unchanged (112 pass, 1 slow deselected)
|
||
- [x] No new binary artifacts committed
|
||
|
||
## Files Changed
|
||
|
||
| File | Change |
|
||
|------|--------|
|
||
| `scripts/smoke_2048.py` | New — main smoke script |
|
||
| `tests/test_smoke_2048.py` | New — pytest smoke + wiring tests |
|
||
| `README.md` | Add "2048 smoke test" subsection |
|
||
|
||
## Unresolved Questions
|
||
|
||
- Exact peak memory at 2048 unknown until first run — might be tight on 8GB machines?
|
||
- Include `--compile` timing breakdown (warmup vs inference) or just total?
|