fix(split-mtp): write the drafter atomically, refuse symlinked outputs
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
Found during the security review of PR #1778; the behavior predates that PR.
## Problem
`split-mtp` checks its output directory once, then, after the shard load and the split, writes `model.safetensors`, `config.json` and up to five companion files (`COMPANION_FILES`: tokenizer, chat template, generation config) into it.
1. The writes follow symlinks: `serialize_to_file`, `std::fs::write` and `std::fs::copy` open the destination with truncation. The guard never looks at the companion names and tests the others with `Path::exists()`, which reports a dangling symlink as absent. In a directory someone else created, a planted `tokenizer.json` symlink to `~/.ssh/authorized_keys` passes the guard without `--force`, and the copy truncates the target. A dangling `config.json` link makes the write create a file wherever it points.
2. The check races the write: a concurrent run passes the same guard, and the two runs' files interleave.
## Evidence
- Writes in `split_mtp_dir`, `src/lib/mlxcel-surgery/src/ops/split_mtp.rs:612-629` (`write_safetensors` at `:614`, reaching `serialize_to_file` at `:522`; `std::fs::write` at `:617`; `std::fs::copy` at `:624`); `COMPANION_FILES` at `:81-87`.
- Guard in `src/commands/split_mtp.rs`: `prepare_output_dir` `:207-237`; `existing_checkpoint_marker` `:249-266` checks only `model.safetensors`, `model.safetensors.index.json` and `config.json` (with `exists()`, `:255`) plus foreign `*.safetensors`. `run_split_mtp` runs the guard at `:77` and the writer at `:98`.
## Fix
This design is decided (maintainer decision, 2026-09-11), including the behavior change for non-empty output directories. Implement it as written; a deviation needs a comment on this issue first.
- `split_mtp_dir` writes every file with `create_new` into a fresh staging directory made with `create_dir` in the output's parent (same filesystem), then commits. Pass the `--force` decision in through `SplitMtpOptions` (default `false`).
- Absent or empty output: one `rename(staging, output)`. `rename` fails on a non-empty target, so the later of two concurrent runs fails instead of interleaving.
- Non-empty output: refuse without `--force`, naming one file the directory holds. This is a deliberate behavior change: today an unrelated non-empty directory is accepted. With `--force`, keep `prepare_output_dir`'s foreign-weights refusal (#1763), refuse if any of the seven destination names is a symlink (`std::fs::symlink_metadata`), then `rename` each staged file in; `rename` replaces the entry and never follows it.
- Remove the staging directory on every failure path; switch `existing_checkpoint_marker` to `symlink_metadata`.
## Acceptance criteria
- [ ] A planted `tokenizer.json` symlink and a dangling `config.json` symlink are refused with and without `--force`, and the link target is left unchanged or never created (tests).
- [ ] Committing onto an output another run filled after staging fails and leaves that run's files intact (test).
- [ ] A non-empty output directory is refused without `--force`, naming a file it holds, and an empty existing directory is accepted and replaced in one `rename` (tests).
- [ ] The `--force` help text (`src/commands/split_mtp.rs:64`) states the new rule, and the PR body lists the behavior change.
- [ ] `--q-bits 4` on `models/mlx/glm-4.7-flash-bf16` still writes `model.safetensors` and `config.json` byte-identical to a build of current main, as PR #1778 recorded.
## Verification
Workspace gate (`cargo test --workspace --profile test-fast --features metal,accelerate`, clippy, `cargo fmt --all -- --check`), then, with `MAIN_OUT` written by a main build:
```bash
./target/release/mlxcel split-mtp -m models/mlx/glm-4.7-flash-bf16 -o "$NEW_OUT" --q-bits 4
cmp "$NEW_OUT/model.safetensors" "$MAIN_OUT/model.safetensors" && cmp "$NEW_OUT/config.json" "$MAIN_OUT/config.json"
```
Contributor guide
Research direction
Start at src/commands/split_mtp.rs, especially run_split_mtp, prepare_output_dir, existing_checkpoint_marker, and the --force help text, then trace split_mtp_dir and write_safetensors in src/lib/mlxcel-surgery/src/ops/split_mtp.rs. Run the targeted split-mtp tests first and add coverage for symlinks, concurrent staging, empty and non-empty outputs, and forced commits. Done means the listed acceptance tests pass, output bytes remain compatible, and the workspace gate and verification commands succeed.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- cli, security
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 45/100