diff --git a/TECHNICAL_REPORTS/pr-2277-atomic-split-mtp-output.en.md b/TECHNICAL_REPORTS/pr-2277-atomic-split-mtp-output.en.md new file mode 100644 index 000000000..49e89ad45 --- /dev/null +++ b/TECHNICAL_REPORTS/pr-2277-atomic-split-mtp-output.en.md @@ -0,0 +1,45 @@ +# PR #2277: Atomic and exclusive `split-mtp` output + +**Date**: 2026-10-12 +**Status**: Implemented, reviewed, and verified with the original GLM-4.7-Flash checkpoint +**Risk**: Medium (checkpoint output commit semantics and a byte-compatible safetensors writer) + +## Summary + +PR #2277 prevents `mlxcel split-mtp` from exposing partial checkpoints or following output-file symlinks. The command now builds the complete drafter in a private sibling staging directory, creates every staged file exclusively, and publishes a new or empty output with one directory rename. A non-empty output is refused unless `--force` is present. + +Forced replacement remains deliberately file-by-file because a populated directory cannot be atomically replaced portably. Before replacement, the command checks the destination again, refuses foreign safetensors files and symlinked drafter paths, and removes only a stale safetensors index plus the known drafter files it replaces. + +## Problem + +The previous directory writer created the destination first and wrote `model.safetensors`, `config.json`, and companion files directly into it. Serialization, JSON writing, and copying used APIs that truncate existing files and follow symlinks. An interrupted conversion could therefore leave a directory that looked like a checkpoint but mixed incomplete and old artifacts. A competing writer could also interleave files after the CLI preflight check. + +The output guard recognized checkpoint-shaped directories but allowed unrelated non-empty directories without `--force`. Path spellings such as `.` and `..`, Unix non-UTF-8 names, and an ancestor symlink retarget during conversion also required explicit handling once staging was introduced. + +## Design + +- A uniquely named staging directory is created beside the resolved output with `create_dir`; on Unix its mode is fixed to `0700`. An RAII guard removes it after every error. +- Safetensors, config, and companion files use `create_new`. The safetensors writer streams its header and tensor payloads to the exclusive file, avoiding a second multi-gigabyte serialization buffer. +- The writer reproduces safetensors 0.7 canonical dtype/name ordering, offsets, metadata, and eight-byte header padding. A mixed-dtype, multi-name regression compares the result byte for byte with `safetensors::serialize_to_file`. +- An absent or empty output is committed with one directory rename. If another writer creates a non-empty output first, the loser fails and its staging directory is removed without changing the winner. +- With `--force`, the commit path rechecks foreign weights, all seven drafter destinations, destination types, and the stale index before renaming staged files individually. +- Path resolution preserves Unix non-UTF-8 names and normalizes unresolved components without creating them. The first resolved target is stored in the staging guard; commit rejects a changed resolution and writes only to that stored target. The source-checkpoint alias guard runs before conversion and again against the stored target. + +## User-visible behavior + +Any non-empty `--output` now requires `--force`, including a directory that contains only an unrelated file. The CLI help explains that foreign safetensors files and symlinked drafter files are refused even with force. Fresh and empty output directories keep the same successful result, while `.` and `..` output spellings and Unix non-UTF-8 names remain supported. + +## Verification + +- Focused surgery unit suite: 26 passed. Coverage includes private permissions, cleanup, a concurrent winner, forced replacement, dangling and present symlinks, ancestor-symlink retargeting, source aliases through unresolved `..`, special path spellings, non-UTF-8 names, exclusive file creation, and serializer equivalence. +- Focused CUDA-feature CLI suite: 10 passed. It covers the default non-empty refusal, force rules, foreign weights, symlink handling, source alias preflight, and help text. +- Formatting and patch whitespace checks passed. The focused Clippy result is recorded in the PR test plan. +- Issue contract checks passed for versions, kernel dtype, port dispatch, Llama compatibility, ROCm overlay, Python 3.12/3.13, and WebUI behavior. +The real-checkpoint acceptance used `zai-org/GLM-4.7-Flash` revision `7dd20894a642a0aa287e9827cb1a1f7f91386b67`. Only the three shards containing the 212 MTP source tensors were needed. Both the baseline at `190b83c3` and the final candidate ran `split-mtp --q-bits 4`, producing 54 output tensors, 15 quantized tensors, and a 0.72 GB payload. `model.safetensors` and `config.json` matched byte for byte: + +- weights SHA-256: `84018ae89ff88bf727f2552523748a21acb4e824d7f22bcf77d8a439d5f4dd61` +- config SHA-256: `7218072d09e5bb06e3eb005b2481ea0b1da08e7e7cd9ad59c69dbd132a754c24` + +## Review findings + +Review added three path regressions before submission. Staging names no longer require UTF-8 and correctly resolve `.` and `..`. The staging guard binds commit to its first resolved target and rejects an ancestor symlink that is retargeted during conversion. Unresolved paths such as `model/new/..` are recognized as source aliases without creating `new` or modifying any source file. The follow-up security review found no remaining issue in scope. diff --git a/TECHNICAL_REPORTS/pr-2277-atomic-split-mtp-output.ko.md b/TECHNICAL_REPORTS/pr-2277-atomic-split-mtp-output.ko.md new file mode 100644 index 000000000..cb598b9ce --- /dev/null +++ b/TECHNICAL_REPORTS/pr-2277-atomic-split-mtp-output.ko.md @@ -0,0 +1,45 @@ +# PR #2277: `split-mtp` 출력을 원자적·배타적으로 커밋하기 + +**날짜**: 2026-10-12 +**상태**: 구현 및 리뷰 완료, 원본 GLM-4.7-Flash 체크포인트로 검증 완료 +**위험도**: 중간 (체크포인트 출력 커밋 방식과 바이트 호환 safetensors writer 변경) + +## 요약 + +PR #2277은 `mlxcel split-mtp`가 일부만 작성된 체크포인트를 노출하거나 출력 파일 심볼릭 링크를 따라가는 문제를 막습니다. 이제 명령은 출력과 같은 부모 아래의 비공개 staging 디렉터리에서 drafter 전체를 만들고, 모든 staging 파일을 배타적으로 생성한 뒤, 새 출력이나 빈 출력은 디렉터리 rename 한 번으로 게시합니다. 비어 있지 않은 출력에는 `--force`가 필요합니다. + +내용이 있는 디렉터리는 이식성 있게 통째로 원자 교체할 수 없으므로 force 교체는 의도적으로 파일 단위로 유지합니다. 교체 직전에 목적지를 다시 검사하고, 외부 safetensors 파일과 심볼릭 링크인 drafter 경로를 거부하며, 오래된 safetensors index와 교체 대상인 알려진 drafter 파일만 다룹니다. + +## 문제 + +기존 디렉터리 writer는 먼저 목적지 디렉터리를 만든 뒤 `model.safetensors`, `config.json`, companion 파일을 그 안에 바로 썼습니다. 직렬화, JSON 기록, 파일 복사에 기존 파일을 truncate하고 심볼릭 링크를 따라가는 API를 사용했습니다. 따라서 변환이 중단되면 불완전한 새 파일과 이전 파일이 섞였지만 체크포인트처럼 보이는 디렉터리가 남을 수 있었습니다. CLI 사전 검사 뒤에 다른 writer가 끼어들어 파일을 섞는 것도 가능했습니다. + +출력 guard는 체크포인트 형태의 디렉터리는 인식했지만, 관련 없는 파일만 있는 비어 있지 않은 디렉터리는 `--force` 없이 허용했습니다. Staging을 도입하면서 `.`, `..`, Unix 비 UTF-8 이름, 변환 도중 조상 심볼릭 링크가 다른 곳을 가리키는 경우도 명시적으로 처리해야 했습니다. + +## 설계 + +- Resolve한 출력과 같은 부모에 고유한 staging 디렉터리를 `create_dir`로 만듭니다. Unix에서는 mode를 `0700`으로 고정합니다. RAII guard가 모든 오류 뒤 staging을 제거합니다. +- Safetensors, config, companion 파일은 `create_new`를 사용합니다. Safetensors writer는 header와 tensor payload를 배타적으로 생성한 파일에 streaming하여 수 GB 크기의 두 번째 직렬화 버퍼를 만들지 않습니다. +- Writer는 safetensors 0.7의 dtype/name 정렬, offset, metadata, 8바이트 header padding을 그대로 재현합니다. 여러 이름과 f32, bf16, u32를 섞은 회귀 테스트가 `safetensors::serialize_to_file` 결과와 바이트 단위로 비교합니다. +- 출력이 없거나 비어 있으면 디렉터리 rename 한 번으로 커밋합니다. 그 사이 다른 writer가 비어 있지 않은 출력을 먼저 만들면 후발 writer는 실패하고 winner를 건드리지 않은 채 staging을 제거합니다. +- `--force`에서는 staged 파일을 개별 rename하기 전에 외부 weight, 7개 drafter 목적지, 목적지 타입, 오래된 index를 다시 검사합니다. +- 경로 해석은 Unix 비 UTF-8 이름을 보존하고 존재하지 않는 component를 만들지 않은 채 정규화합니다. 최초 resolve 결과를 staging guard에 저장하고, 커밋 시 resolve 결과가 달라지면 거부하며, 실제 기록은 저장한 target에만 수행합니다. 소스 체크포인트 alias 검사는 변환 전과 저장한 target을 대상으로 한 번 더 실행합니다. + +## 사용자에게 보이는 변경 + +관련 없는 파일 하나만 있는 경우를 포함해 비어 있지 않은 모든 `--output`에는 이제 `--force`가 필요합니다. CLI help는 force를 사용해도 외부 safetensors 파일과 심볼릭 링크인 drafter 파일을 거부한다고 설명합니다. 새 출력과 빈 출력의 성공 결과는 같고, `.`, `..` 표기와 Unix 비 UTF-8 출력 이름도 계속 지원합니다. + +## 검증 + +- Surgery 집중 단위 테스트: 26개 통과. 비공개 권한, cleanup, 동시 winner, force 교체, 정상·dangling 심볼릭 링크, 조상 심볼릭 링크 retarget, resolve되지 않은 `..`를 통한 source alias, 특수 경로 표기, 비 UTF-8 이름, 배타적 파일 생성, serializer 동등성을 포함합니다. +- CUDA feature CLI 집중 테스트: 10개 통과. 기본 non-empty 거부, force 규칙, 외부 weight, 심볼릭 링크, source alias 사전 검사, help 문구를 확인합니다. +- 포맷과 patch whitespace 검사가 통과했습니다. 집중 Clippy 결과는 PR test plan에 기록했습니다. +- 버전, kernel dtype, port dispatch, Llama 호환성, ROCm overlay, Python 3.12/3.13, WebUI 동작에 대한 이슈 contract 검사가 통과했습니다. +실제 체크포인트 검증에는 `zai-org/GLM-4.7-Flash` revision `7dd20894a642a0aa287e9827cb1a1f7f91386b67`을 사용했습니다. MTP 원본 tensor 212개가 든 shard 3개만 필요했습니다. `190b83c3` baseline과 최종 candidate 모두 `split-mtp --q-bits 4`를 실행해 출력 tensor 54개, quantized tensor 15개, 0.72 GB payload를 만들었습니다. `model.safetensors`와 `config.json`은 바이트 단위로 일치했습니다. + +- weights SHA-256: `84018ae89ff88bf727f2552523748a21acb4e824d7f22bcf77d8a439d5f4dd61` +- config SHA-256: `7218072d09e5bb06e3eb005b2481ea0b1da08e7e7cd9ad59c69dbd132a754c24` + +## 리뷰 결과 + +제출 전 리뷰에서 경로 관련 회귀 세 가지를 보완했습니다. Staging 이름은 더 이상 UTF-8을 요구하지 않고 `.`과 `..`를 올바르게 resolve합니다. Staging guard는 최초 resolve target에 커밋을 고정하고, 변환 도중 조상 심볼릭 링크가 다른 곳을 가리키면 이를 거부합니다. `model/new/..`처럼 아직 resolve되지 않은 경로도 `new`를 만들거나 소스 파일을 수정하지 않은 채 source alias로 인식합니다. 후속 보안 리뷰에서 이 범위의 추가 문제는 발견되지 않았습니다. diff --git a/src/commands/split_mtp.rs b/src/commands/split_mtp.rs index 3e4a7e205..ed68d78b9 100644 --- a/src/commands/split_mtp.rs +++ b/src/commands/split_mtp.rs @@ -61,7 +61,9 @@ pub(crate) struct SplitMtpArgs { #[arg(long, default_value_t = 64, value_name = "N")] pub(crate) q_group_size: i32, - /// Overwrite an output directory that already holds a checkpoint. + /// A non-empty output directory is refused unless this flag is set. + /// Foreign safetensors weights and symlinked drafter files are always + /// refused. #[arg(long)] pub(crate) force: bool, } @@ -85,6 +87,7 @@ pub(crate) fn run_split_mtp(args: SplitMtpArgs) -> Result<()> { block_size: args.block_size, q_bits: args.q_bits, q_group_size: args.q_group_size, + force: args.force, }; println!( "Splitting the MTP block out of {} into {}{}", @@ -189,18 +192,13 @@ fn preflight_checks( /// Guard `args.output` before [`split_mtp_dir`] writes into it. /// -/// No existing marker: `Ok(())`, nothing to guard. A marker with `force` -/// unset: today's `already holds a checkpoint (...); pass --force to -/// overwrite` refusal, unchanged. A marker with `force` set and at least one -/// foreign safetensors file present ([`is_foreign_safetensors_file`]): a new -/// refusal naming one such file, because `--force` overwrites the drafter's -/// own `model.safetensors` and `config.json` but does not promise to delete -/// weight files, and the file it would otherwise leave orphaned may be the -/// only surviving copy (issue #1763). A marker with `force` set and no -/// foreign files: remove a stale `model.safetensors.index.json` as before -/// and return `Ok(())`, since `collect_shard_paths` prefers an index over a -/// bare `model.safetensors` and a stale one left beside the drafter's -/// single-file output would send the loader to the victim's shards. +/// An absent or empty directory is accepted. Any non-empty directory without +/// `force` is refused, including one that holds only unrelated files. With +/// `force`, foreign safetensors files remain protected (issue #1763) and all +/// seven drafter destinations are checked with `symlink_metadata`, so a +/// dangling link is refused as well. This function is intentionally +/// non-mutating; [`split_mtp_dir`] stages the complete drafter and owns the +/// commit-time checks that close races after this early guard. /// /// Used by: [`run_split_mtp`] (after [`preflight_checks`] has already /// refused anything that does not require touching the filesystem). @@ -210,7 +208,7 @@ fn prepare_output_dir(dir: &std::path::Path, force: bool) -> Result<()> { }; if !force { return Err(anyhow!( - "split-mtp: {} already holds a checkpoint ({existing}); pass --force to overwrite", + "split-mtp: {} is not empty ({existing}); pass --force to overwrite drafter files", dir.display() )); } @@ -228,22 +226,31 @@ fn prepare_output_dir(dir: &std::path::Path, force: bool) -> Result<()> { dir.display() )); } - let index = dir.join("model.safetensors.index.json"); - if index.exists() { - std::fs::remove_file(&index) - .with_context(|| format!("split-mtp: failed to remove stale {}", index.display()))?; + for name in [ + "model.safetensors", + "config.json", + "tokenizer.json", + "tokenizer_config.json", + "chat_template.jinja", + "generation_config.json", + "special_tokens_map.json", + ] { + let path = dir.join(name); + if std::fs::symlink_metadata(&path).is_ok_and(|metadata| metadata.file_type().is_symlink()) + { + return Err(anyhow!( + "split-mtp: refusing symlinked output file {}", + path.display() + )); + } } Ok(()) } -/// Name of the first artifact that makes `dir` look like an existing -/// checkpoint, or `None` when writing there would clobber nothing. -/// -/// `split_mtp_dir` writes `model.safetensors` and `config.json` and copies -/// five tokenizer files over whatever is present, so the guard has to screen -/// on more than the single-file weight name: a sharded checkpoint has an -/// index and shards instead, and would lose its config and tokenizer while -/// keeping shards nothing can load. +/// Name of one entry that makes `dir` non-empty, or `None` when writing there +/// would clobber nothing. Known drafter and checkpoint names are checked first +/// with `symlink_metadata` so dangling links are visible; remaining entries +/// are returned in filename order. /// /// Used by: [`run_split_mtp`] (through [`prepare_output_dir`]). fn existing_checkpoint_marker(dir: &std::path::Path) -> Option { @@ -251,18 +258,23 @@ fn existing_checkpoint_marker(dir: &std::path::Path) -> Option { "model.safetensors", "model.safetensors.index.json", "config.json", + "tokenizer.json", + "tokenizer_config.json", + "chat_template.jinja", + "generation_config.json", + "special_tokens_map.json", ] { - if dir.join(name).exists() { + if std::fs::symlink_metadata(dir.join(name)).is_ok() { return Some(name.to_string()); } } - let shard = std::fs::read_dir(dir).ok()?.flatten().find(|entry| { - entry - .file_name() - .to_str() - .is_some_and(is_foreign_safetensors_file) - })?; - Some(shard.file_name().to_string_lossy().into_owned()) + let mut names: Vec = std::fs::read_dir(dir) + .ok()? + .flatten() + .map(|entry| entry.file_name().to_string_lossy().into_owned()) + .collect(); + names.sort(); + names.into_iter().next() } #[cfg(test)] @@ -297,11 +309,14 @@ mod tests { } #[test] - fn marker_is_none_for_an_empty_or_unrelated_directory() { + fn marker_is_none_only_for_an_empty_or_missing_directory() { let dir = tempfile::tempdir().expect("tempdir"); assert_eq!(existing_checkpoint_marker(dir.path()), None); std::fs::write(dir.path().join("README.md"), b"x").expect("write"); - assert_eq!(existing_checkpoint_marker(dir.path()), None); + assert_eq!( + existing_checkpoint_marker(dir.path()).as_deref(), + Some("README.md") + ); assert_eq!( existing_checkpoint_marker(&dir.path().join("missing")), None, @@ -334,11 +349,10 @@ mod tests { assert!(index.exists(), "the index must not be deleted"); } - /// A stale index with no surviving shards is still cleared under - /// `--force`, as before: there is nothing left for it to misdirect the - /// loader toward. + /// The early guard is non-mutating. The staged writer removes the stale + /// index only when the replacement drafter is ready to commit. #[test] - fn force_clears_a_stale_index_when_no_shards_remain() { + fn force_accepts_a_stale_index_without_mutating_the_output() { let dir = tempfile::tempdir().expect("tempdir"); let weights = dir.path().join("model.safetensors"); let index = dir.path().join("model.safetensors.index.json"); @@ -347,10 +361,60 @@ mod tests { prepare_output_dir(dir.path(), true).expect("must be accepted"); - assert!(!index.exists(), "the stale index must be removed"); + assert!(index.exists(), "preflight must leave the stale index alone"); assert!(weights.exists(), "model.safetensors must be left alone"); } + #[test] + fn non_empty_output_is_refused_without_force_and_names_an_entry() { + let dir = tempfile::tempdir().expect("tempdir"); + let note = dir.path().join("README.md"); + std::fs::write(¬e, b"keep").expect("write"); + + let err = prepare_output_dir(dir.path(), false).expect_err("must refuse"); + assert!(err.to_string().contains("README.md"), "{err}"); + assert_eq!(std::fs::read(¬e).expect("read"), b"keep"); + } + + #[cfg(unix)] + #[test] + fn force_refuses_present_and_dangling_drafter_symlinks() { + use std::os::unix::fs::symlink; + + for name in ["tokenizer.json", "config.json"] { + let dir = tempfile::tempdir().expect("tempdir"); + let target = dir.path().join(format!("{name}.target")); + if name == "tokenizer.json" { + std::fs::write(&target, b"untouched").expect("write target"); + } + symlink(&target, dir.path().join(name)).expect("symlink"); + + let err = prepare_output_dir(dir.path(), true).expect_err("must refuse"); + assert!(err.to_string().contains(name), "{err}"); + if name == "tokenizer.json" { + assert_eq!(std::fs::read(&target).expect("read target"), b"untouched"); + } else { + assert!(!target.exists(), "dangling target must not be created"); + } + } + } + + #[test] + fn force_help_describes_the_non_empty_and_symlink_rules() { + use clap::CommandFactory; + + let mut command = crate::Cli::command(); + let split_mtp = command + .find_subcommand_mut("split-mtp") + .expect("split-mtp command"); + let help = split_mtp.render_long_help().to_string(); + assert!( + help.contains("A non-empty output directory is refused unless this flag is set"), + "{help}" + ); + assert!(help.contains("symlinked drafter files"), "{help}"); + } + /// `--force` must also refuse a directory holding a foreign safetensors /// file that is not a `model-*`-named shard: `consolidated.safetensors` /// passes the loader's `glob_safetensors` glob just as readily as a diff --git a/src/lib/mlxcel-surgery/src/ops/split_mtp.rs b/src/lib/mlxcel-surgery/src/ops/split_mtp.rs index 9fa3fa203..57c12ca6e 100644 --- a/src/lib/mlxcel-surgery/src/ops/split_mtp.rs +++ b/src/lib/mlxcel-surgery/src/ops/split_mtp.rs @@ -58,8 +58,7 @@ //! for its own layers, so `TransformerBlock::from_weights` loads the block //! unchanged. -use std::collections::HashMap; -use std::path::{Path, PathBuf}; +use std::path::PathBuf; use anyhow::anyhow; use mlxcel_core::mla::{KvBProjGeometry, decompose_kv_b_proj}; @@ -98,6 +97,9 @@ pub struct SplitMtpOptions { pub q_bits: Option, /// Affine quantization group size. Only read when `q_bits` is set. pub q_group_size: i32, + /// Replace drafter files in a non-empty output directory. Foreign + /// safetensors files and symlinked drafter paths are always refused. + pub force: bool, } impl Default for SplitMtpOptions { @@ -106,6 +108,7 @@ impl Default for SplitMtpOptions { block_size: None, q_bits: None, q_group_size: 64, + force: false, } } } @@ -446,198 +449,12 @@ pub fn split_mtp( }) } -/// An evaluated tensor held as host bytes for the safetensors writer. -struct OwnedTensor { - dtype: safetensors::Dtype, - shape: Vec, - data: Vec, -} - -impl safetensors::View for &OwnedTensor { - fn dtype(&self) -> safetensors::Dtype { - self.dtype - } - fn shape(&self) -> &[usize] { - &self.shape - } - fn data(&self) -> std::borrow::Cow<'_, [u8]> { - std::borrow::Cow::Borrowed(&self.data) - } - fn data_len(&self) -> usize { - self.data.len() - } -} - -fn safetensors_dtype(mlx_dtype: i32) -> Result { - use mlxcel_core::dtype as d; - use safetensors::Dtype; - Ok(match mlx_dtype { - d::BOOL => Dtype::BOOL, - d::UINT8 => Dtype::U8, - d::UINT16 => Dtype::U16, - d::UINT32 => Dtype::U32, - d::UINT64 => Dtype::U64, - d::INT8 => Dtype::I8, - d::INT16 => Dtype::I16, - d::INT32 => Dtype::I32, - d::INT64 => Dtype::I64, - d::FLOAT16 => Dtype::F16, - d::FLOAT32 => Dtype::F32, - d::FLOAT64 => Dtype::F64, - d::BFLOAT16 => Dtype::BF16, - other => { - return Err( - anyhow!("split-mtp: MLX dtype {other} has no safetensors equivalent").into(), - ); - } - }) -} - -/// Materialize `weights` and write them as one `model.safetensors` with the -/// `{"format": "mlx"}` metadata every mlx-community checkpoint carries. -/// -/// Consumes the map so each array is released as soon as its bytes are on -/// the host; the drafter's bf16 payload is a few GB and holding both copies -/// would double the peak. -pub fn write_safetensors(weights: WeightMap, path: &Path) -> Result { - let mut owned: Vec<(String, OwnedTensor)> = Vec::with_capacity(weights.len()); - let mut bytes_written = 0u64; - let mut entries: Vec<(String, UniquePtr)> = weights.into_iter().collect(); - entries.sort_by(|a, b| a.0.cmp(&b.0)); - for (key, array) in entries { - let dtype = safetensors_dtype(mlxcel_core::array_dtype(&array))?; - let shape: Vec = mlxcel_core::array_shape(&array) - .iter() - .map(|&d| d as usize) - .collect(); - mlxcel_core::eval(&array); - let data = mlxcel_core::array_to_raw_bytes(&array); - drop(array); - bytes_written += data.len() as u64; - owned.push((key, OwnedTensor { dtype, shape, data })); - } - let views: Vec<(&str, &OwnedTensor)> = owned.iter().map(|(k, t)| (k.as_str(), t)).collect(); - let mut metadata = HashMap::new(); - metadata.insert("format".to_string(), "mlx".to_string()); - safetensors::serialize_to_file(views, Some(metadata), path) - .map_err(|e| anyhow!("split-mtp: writing {}: {e}", path.display()))?; - Ok(bytes_written) -} - -/// Whether the safetensors index names any `model.layers.{layer}.*` tensor. -/// -/// One JSON parse, no shard opened. `Ok(true)` for a directory without an -/// index: the filtered load below then answers the question from the shard -/// headers and [`split_mtp`] reports the same refusal on an empty result. -fn index_names_nextn_layer(model_dir: &Path, layer: usize) -> Result { - let index_path = model_dir.join("model.safetensors.index.json"); - if !index_path.exists() { - return Ok(true); - } - let prefix = format!("model.layers.{layer}."); - let raw = std::fs::read_to_string(&index_path)?; - let value: Value = serde_json::from_str(&raw) - .map_err(|e| anyhow!("split-mtp: parsing {}: {e}", index_path.display()))?; - let map = value - .get("weight_map") - .and_then(Value::as_object) - .ok_or_else(|| anyhow!("{}: missing weight_map object", index_path.display()))?; - Ok(map.keys().any(|k| k.starts_with(&prefix))) -} - -/// Refuse when `output_dir` resolves to the same directory as `model_dir`. -/// -/// A sharded raw checkpoint carries no `model.safetensors`, so the CLI's -/// overwrite guard (`prepare_output_dir` in the `mlxcel` binary crate) does -/// not fire on it, and the writes [`split_mtp_dir`] performs would replace -/// the source `config.json` with the drafter's and truncate every companion -/// file: `std::fs::copy` on a path to itself reports `Ok(0)` after opening -/// the destination with `O_TRUNC`. This has to run before any mutation, -/// including the CLI's own overwrite guard, which cannot tell "the shard I'm -/// looking at belongs to the source" from "it belongs to a stray previous -/// output" (PR #1778 review). -/// -/// Used by: [`split_mtp_dir`], and -/// `crate::commands::split_mtp::run_split_mtp` (through its -/// `preflight_checks`) in the consuming crate -pub fn refuse_output_is_source(model_dir: &Path, output_dir: &Path) -> Result<(), SurgeryError> { - if output_dir.exists() - && std::fs::canonicalize(model_dir)? == std::fs::canonicalize(output_dir)? - { - return Err(anyhow!( - "split-mtp: --output {} is the source checkpoint; write the drafter to a separate \ - directory, otherwise the source config.json is overwritten and its tokenizer \ - files are truncated", - output_dir.display() - ) - .into()); - } - Ok(()) -} - -/// Split the nextn block out of `model_dir` into `output_dir`. -/// -/// Reads `config.json`, refuses a checkpoint whose index names no nextn -/// tensor (the community conversions), loads only the shards that hold the -/// block, runs [`split_mtp`], writes `model.safetensors` and `config.json`, -/// and copies the tokenizer files alongside. -pub fn split_mtp_dir( - model_dir: &Path, - output_dir: &Path, - opts: &SplitMtpOptions, -) -> Result { - refuse_output_is_source(model_dir, output_dir)?; - let config_path = model_dir.join("config.json"); - let config: Value = serde_json::from_str(&std::fs::read_to_string(&config_path)?) - .map_err(|e| anyhow!("split-mtp: parsing {}: {e}", config_path.display()))?; - let source_layer = nextn_layer_index(&config)?; - - if !index_names_nextn_layer(model_dir, source_layer)? { - return Err(anyhow!( - "split-mtp: {} names no `model.layers.{source_layer}.*` tensor; this checkpoint was \ - converted without its MTP layer; use the raw zai-org checkpoint", - model_dir.display() - ) - .into()); - } - - let layer_prefix = format!("model.layers.{source_layer}."); - let weights = mlxcel_core::weights::load_weights_from_dir_index_filtered(model_dir, |name| { - name.starts_with(&layer_prefix) - }) - .map_err(|e| anyhow!("split-mtp: {e}"))?; - - let result = split_mtp(weights, &config, opts)?; - - std::fs::create_dir_all(output_dir)?; - let tensors = result.weights.len(); - let bytes_written = write_safetensors(result.weights, &output_dir.join("model.safetensors"))?; - let config_text = serde_json::to_string_pretty(&result.config) - .map_err(|e| anyhow!("split-mtp: serializing config.json: {e}"))?; - std::fs::write(output_dir.join("config.json"), config_text)?; - - let mut copied_files = Vec::new(); - let mut missing_files = Vec::new(); - for name in COMPANION_FILES { - let src = model_dir.join(name); - if src.is_file() { - std::fs::copy(&src, output_dir.join(name))?; - copied_files.push((*name).to_string()); - } else if *name == "tokenizer.json" || *name == "tokenizer_config.json" { - missing_files.push((*name).to_string()); - } - } - - Ok(SplitMtpReport { - output_dir: output_dir.to_path_buf(), - source_layer, - tensors, - quantized_tensors: result.quantized_tensors, - bytes_written, - copied_files, - missing_files, - }) -} +#[path = "split_mtp_dir.rs"] +mod directory; +pub use directory::{refuse_output_is_source, split_mtp_dir}; +#[path = "split_mtp_output.rs"] +mod output; +pub use output::write_safetensors; #[cfg(test)] #[path = "split_mtp_tests.rs"] diff --git a/src/lib/mlxcel-surgery/src/ops/split_mtp_dir.rs b/src/lib/mlxcel-surgery/src/ops/split_mtp_dir.rs new file mode 100644 index 000000000..a2cf01672 --- /dev/null +++ b/src/lib/mlxcel-surgery/src/ops/split_mtp_dir.rs @@ -0,0 +1,91 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Checkpoint directory driver for `split-mtp`. + +use std::path::Path; + +use anyhow::anyhow; +use serde_json::Value; + +use super::{SplitMtpOptions, SplitMtpReport, nextn_layer_index, output, split_mtp}; +use crate::SurgeryError; + +/// Whether the safetensors index names any `model.layers.{layer}.*` tensor. +/// +/// One JSON parse, no shard opened. `Ok(true)` for a directory without an +/// index: the filtered load below then answers the question from the shard +/// headers and [`split_mtp`] reports the same refusal on an empty result. +fn index_names_nextn_layer(model_dir: &Path, layer: usize) -> Result { + let index_path = model_dir.join("model.safetensors.index.json"); + if !index_path.exists() { + return Ok(true); + } + let prefix = format!("model.layers.{layer}."); + let raw = std::fs::read_to_string(&index_path)?; + let value: Value = serde_json::from_str(&raw) + .map_err(|error| anyhow!("split-mtp: parsing {}: {error}", index_path.display()))?; + let map = value + .get("weight_map") + .and_then(Value::as_object) + .ok_or_else(|| anyhow!("{}: missing weight_map object", index_path.display()))?; + Ok(map.keys().any(|key| key.starts_with(&prefix))) +} + +/// Refuse when `output_dir` resolves to the same directory as `model_dir`. +/// +/// A sharded raw checkpoint carries no `model.safetensors`, so the CLI's +/// overwrite guard does not fire on it. The refusal runs before any mutation, +/// since the CLI cannot distinguish a source shard from a foreign output. +pub fn refuse_output_is_source(model_dir: &Path, output_dir: &Path) -> Result<(), SurgeryError> { + if std::fs::canonicalize(model_dir)? == output::resolve_output_path(output_dir)? { + return Err(output::source_alias_error(model_dir)); + } + Ok(()) +} + +/// Split the nextn block out of `model_dir` into `output_dir`. +/// +/// Reads `config.json`, refuses a checkpoint whose index names no nextn +/// tensor, loads only the shards that hold the block, performs the transform, +/// and sends the complete result through the exclusive staging writer. +pub fn split_mtp_dir( + model_dir: &Path, + output_dir: &Path, + opts: &SplitMtpOptions, +) -> Result { + refuse_output_is_source(model_dir, output_dir)?; + let config_path = model_dir.join("config.json"); + let config: Value = serde_json::from_str(&std::fs::read_to_string(&config_path)?) + .map_err(|error| anyhow!("split-mtp: parsing {}: {error}", config_path.display()))?; + let source_layer = nextn_layer_index(&config)?; + + if !index_names_nextn_layer(model_dir, source_layer)? { + return Err(anyhow!( + "split-mtp: {} names no `model.layers.{source_layer}.*` tensor; this checkpoint was \ + converted without its MTP layer; use the raw zai-org checkpoint", + model_dir.display() + ) + .into()); + } + + let layer_prefix = format!("model.layers.{source_layer}."); + let weights = mlxcel_core::weights::load_weights_from_dir_index_filtered(model_dir, |name| { + name.starts_with(&layer_prefix) + }) + .map_err(|error| anyhow!("split-mtp: {error}"))?; + let result = split_mtp(weights, &config, opts)?; + + output::write_split_output(model_dir, output_dir, result, source_layer, opts.force) +} diff --git a/src/lib/mlxcel-surgery/src/ops/split_mtp_output.rs b/src/lib/mlxcel-surgery/src/ops/split_mtp_output.rs new file mode 100644 index 000000000..9795ffc63 --- /dev/null +++ b/src/lib/mlxcel-surgery/src/ops/split_mtp_output.rs @@ -0,0 +1,486 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Exclusive staging and commit support for `split-mtp` output directories. + +use std::collections::HashMap; +use std::ffi::OsString; +use std::fs::{File, OpenOptions}; +use std::io::{BufWriter, Write}; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU64, Ordering}; + +use anyhow::{Context, anyhow}; +use mlxcel_core::weights::WeightMap; +use mlxcel_core::{MlxArray, UniquePtr}; +use safetensors::tensor::{Metadata, TensorInfo}; + +use super::{COMPANION_FILES, SplitMtpReport, SplitMtpResult}; +use crate::SurgeryError; + +const OUTPUT_FILES: &[&str] = &[ + "model.safetensors", + "config.json", + "tokenizer.json", + "tokenizer_config.json", + "chat_template.jinja", + "generation_config.json", + "special_tokens_map.json", +]; +const STALE_INDEX: &str = "model.safetensors.index.json"; +const MAX_HEADER_SIZE: usize = 100_000_000; +static STAGING_SEQUENCE: AtomicU64 = AtomicU64::new(0); + +struct OwnedTensor { + dtype: safetensors::Dtype, + shape: Vec, + data: Vec, +} + +/// Materialize `weights` and write one exclusive `model.safetensors` with +/// the same byte layout as `safetensors::serialize_to_file`. +/// +/// The crate writer uses `File::create`, which follows links and truncates +/// existing files. This streaming equivalent uses `create_new` and avoids a +/// second multi-gigabyte serialization buffer. +pub fn write_safetensors(weights: WeightMap, path: &Path) -> Result { + let mut owned: Vec<(String, OwnedTensor)> = Vec::with_capacity(weights.len()); + let mut bytes_written = 0u64; + let mut entries: Vec<(String, UniquePtr)> = weights.into_iter().collect(); + entries.sort_by(|a, b| a.0.cmp(&b.0)); + for (key, array) in entries { + let dtype = safetensors_dtype(mlxcel_core::array_dtype(&array))?; + let shape: Vec = mlxcel_core::array_shape(&array) + .iter() + .map(|&dimension| dimension as usize) + .collect(); + mlxcel_core::eval(&array); + let data = mlxcel_core::array_to_raw_bytes(&array); + drop(array); + bytes_written = bytes_written + .checked_add(data.len() as u64) + .ok_or_else(|| anyhow!("split-mtp: safetensors payload size overflow"))?; + owned.push((key, OwnedTensor { dtype, shape, data })); + } + + // Match safetensors 0.7's canonical order: descending dtype alignment, + // then tensor name. Header offsets and payload order depend on it. + owned.sort_by(|(left_name, left), (right_name, right)| { + right.dtype.cmp(&left.dtype).then(left_name.cmp(right_name)) + }); + let mut offset = 0usize; + let mut tensor_info = Vec::with_capacity(owned.len()); + for (name, tensor) in &owned { + let end = offset + .checked_add(tensor.data.len()) + .ok_or_else(|| anyhow!("split-mtp: safetensors payload offset overflow"))?; + tensor_info.push(( + name.clone(), + TensorInfo { + dtype: tensor.dtype, + shape: tensor.shape.clone(), + data_offsets: (offset, end), + }, + )); + offset = end; + } + let mut data_info = HashMap::new(); + data_info.insert("format".to_string(), "mlx".to_string()); + let metadata = Metadata::new(Some(data_info), tensor_info) + .map_err(|error| anyhow!("split-mtp: preparing safetensors metadata: {error}"))?; + let mut header = serde_json::to_vec(&metadata) + .map_err(|error| anyhow!("split-mtp: serializing safetensors metadata: {error}"))?; + let aligned_len = header + .len() + .checked_add(7) + .ok_or_else(|| anyhow!("split-mtp: safetensors header size overflow"))? + / 8 + * 8; + if aligned_len > MAX_HEADER_SIZE { + return Err(anyhow!("split-mtp: safetensors header too large").into()); + } + header.resize(aligned_len, b' '); + + let file = OpenOptions::new() + .write(true) + .create_new(true) + .open(path) + .with_context(|| format!("split-mtp: creating {}", path.display()))?; + let mut writer = BufWriter::new(file); + writer.write_all(&(aligned_len as u64).to_le_bytes())?; + writer.write_all(&header)?; + for (_, tensor) in owned { + writer.write_all(&tensor.data)?; + } + writer.flush()?; + Ok(bytes_written) +} + +fn safetensors_dtype(mlx_dtype: i32) -> Result { + use mlxcel_core::dtype as d; + use safetensors::Dtype; + Ok(match mlx_dtype { + d::BOOL => Dtype::BOOL, + d::UINT8 => Dtype::U8, + d::UINT16 => Dtype::U16, + d::UINT32 => Dtype::U32, + d::UINT64 => Dtype::U64, + d::INT8 => Dtype::I8, + d::INT16 => Dtype::I16, + d::INT32 => Dtype::I32, + d::INT64 => Dtype::I64, + d::FLOAT16 => Dtype::F16, + d::FLOAT32 => Dtype::F32, + d::FLOAT64 => Dtype::F64, + d::BFLOAT16 => Dtype::BF16, + other => { + return Err( + anyhow!("split-mtp: MLX dtype {other} has no safetensors equivalent").into(), + ); + } + }) +} + +pub(super) fn write_split_output( + model_dir: &Path, + output_dir: &Path, + result: SplitMtpResult, + source_layer: usize, + force: bool, +) -> Result { + let staging = StagingDir::create(output_dir)?; + if std::fs::canonicalize(model_dir)? == staging.output_path { + return Err(source_alias_error(model_dir)); + } + let tensors = result.weights.len(); + let bytes_written = + write_safetensors(result.weights, &staging.path().join("model.safetensors"))?; + let config_text = serde_json::to_string_pretty(&result.config) + .map_err(|error| anyhow!("split-mtp: serializing config.json: {error}"))?; + write_new(&staging.path().join("config.json"), config_text.as_bytes())?; + + let mut copied_files = Vec::new(); + let mut missing_files = Vec::new(); + for name in COMPANION_FILES { + let source = model_dir.join(name); + if source.is_file() { + copy_new(&source, &staging.path().join(name))?; + copied_files.push((*name).to_string()); + } else if *name == "tokenizer.json" || *name == "tokenizer_config.json" { + missing_files.push((*name).to_string()); + } + } + + commit_staging_dir(staging, output_dir, force)?; + Ok(SplitMtpReport { + output_dir: output_dir.to_path_buf(), + source_layer, + tensors, + quantized_tensors: result.quantized_tensors, + bytes_written, + copied_files, + missing_files, + }) +} + +fn write_new(path: &Path, contents: &[u8]) -> Result<(), SurgeryError> { + let mut file = OpenOptions::new() + .write(true) + .create_new(true) + .open(path) + .with_context(|| format!("split-mtp: creating {}", path.display()))?; + file.write_all(contents)?; + Ok(()) +} + +fn copy_new(source: &Path, destination: &Path) -> Result<(), SurgeryError> { + let mut source_file = File::open(source)?; + let mut destination_file = OpenOptions::new() + .write(true) + .create_new(true) + .open(destination) + .with_context(|| format!("split-mtp: creating {}", destination.display()))?; + std::io::copy(&mut source_file, &mut destination_file)?; + Ok(()) +} + +pub(super) struct StagingDir { + path: PathBuf, + output_path: PathBuf, + committed: bool, +} + +impl StagingDir { + pub(super) fn create(output_dir: &Path) -> Result { + let resolved_output = resolve_output_path(output_dir)?; + let parent = resolved_output + .parent() + .filter(|path| !path.as_os_str().is_empty()) + .unwrap_or_else(|| Path::new(".")); + std::fs::create_dir_all(parent)?; + let output_name = resolved_output + .file_name() + .ok_or_else(|| anyhow!("split-mtp: output path must name a directory"))?; + + for _ in 0..1000 { + let sequence = STAGING_SEQUENCE.fetch_add(1, Ordering::Relaxed); + let mut staging_name = OsString::from("."); + staging_name.push(output_name); + staging_name.push(format!(".split-mtp-{}-{sequence}.tmp", std::process::id())); + let path = parent.join(staging_name); + match create_private_directory(&path) { + Ok(()) => { + return Ok(Self { + path, + output_path: resolved_output, + committed: false, + }); + } + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => continue, + Err(error) => return Err(error.into()), + } + } + Err(anyhow!( + "split-mtp: could not create a unique staging directory beside {}", + output_dir.display() + ) + .into()) + } + + pub(super) fn path(&self) -> &Path { + &self.path + } +} + +enum UnresolvedComponent { + Normal(OsString), + Parent, +} + +/// Resolve the longest existing prefix, then normalize missing components +/// without creating them. This preserves non-UTF-8 names and resolves `..` +/// relative to an ancestor symlink's target, matching filesystem traversal. +pub(super) fn resolve_output_path(output_dir: &Path) -> Result { + let mut candidate = if output_dir.is_absolute() { + output_dir.to_path_buf() + } else { + std::env::current_dir()?.join(output_dir) + }; + let mut unresolved = Vec::new(); + loop { + match std::fs::symlink_metadata(&candidate) { + Ok(metadata) => { + if unresolved.is_empty() && metadata.file_type().is_symlink() { + return Err(anyhow!( + "split-mtp: output directory {} is a symlink", + output_dir.display() + ) + .into()); + } + let mut resolved = std::fs::canonicalize(&candidate)?; + for component in unresolved.into_iter().rev() { + match component { + UnresolvedComponent::Normal(name) => resolved.push(name), + UnresolvedComponent::Parent => { + resolved.pop(); + } + } + } + return Ok(resolved); + } + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => return Err(error.into()), + } + + let component = candidate + .components() + .next_back() + .ok_or_else(|| anyhow!("split-mtp: output path must name a directory"))?; + match component { + std::path::Component::Normal(name) => { + unresolved.push(UnresolvedComponent::Normal(name.to_os_string())); + } + std::path::Component::ParentDir => unresolved.push(UnresolvedComponent::Parent), + std::path::Component::CurDir => {} + std::path::Component::RootDir | std::path::Component::Prefix(_) => { + return Err(anyhow!("split-mtp: output path must name a directory").into()); + } + } + if !candidate.pop() { + return Err(anyhow!("split-mtp: output path must name a directory").into()); + } + } +} + +fn create_private_directory(path: &Path) -> std::io::Result<()> { + let mut builder = std::fs::DirBuilder::new(); + #[cfg(unix)] + { + use std::os::unix::fs::DirBuilderExt; + + builder.mode(0o700); + } + builder.create(path) +} + +impl Drop for StagingDir { + fn drop(&mut self) { + if !self.committed { + let _ = std::fs::remove_dir_all(&self.path); + } + } +} + +pub(super) fn commit_staging_dir( + mut staging: StagingDir, + output_dir: &Path, + force: bool, +) -> Result<(), SurgeryError> { + let current_output = resolve_output_path(output_dir)?; + if current_output != staging.output_path { + return Err(anyhow!( + "split-mtp: output path {} changed while the drafter was staged", + output_dir.display() + ) + .into()); + } + let output_path = staging.output_path.clone(); + let output_dir = output_path.as_path(); + match first_directory_entry(output_dir)? { + None => { + std::fs::rename(staging.path(), output_dir).with_context(|| { + format!( + "split-mtp: committing staged drafter to {}", + output_dir.display() + ) + })?; + staging.committed = true; + Ok(()) + } + Some(existing) if !force => Err(anyhow!( + "split-mtp: {} is not empty ({existing}); pass --force to overwrite drafter files", + output_dir.display() + ) + .into()), + Some(_) => commit_over_existing(&mut staging, output_dir), + } +} + +pub(super) fn source_alias_error(model_dir: &Path) -> SurgeryError { + anyhow!( + "split-mtp: --output {} is the source checkpoint; write the drafter to a separate \ + directory, otherwise the source config.json is overwritten and its tokenizer \ + files are truncated", + model_dir.display() + ) + .into() +} + +fn first_directory_entry(dir: &Path) -> Result, SurgeryError> { + match std::fs::symlink_metadata(dir) { + Ok(metadata) if metadata.file_type().is_symlink() => { + return Err( + anyhow!("split-mtp: output directory {} is a symlink", dir.display()).into(), + ); + } + Ok(metadata) if !metadata.is_dir() => { + return Err(anyhow!( + "split-mtp: output path {} exists and is not a directory", + dir.display() + ) + .into()); + } + Ok(_) => {} + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(error) => return Err(error.into()), + } + + let mut names = Vec::new(); + for entry in std::fs::read_dir(dir)? { + names.push(entry?.file_name().to_string_lossy().into_owned()); + } + names.sort(); + Ok(names.into_iter().next()) +} + +fn commit_over_existing(staging: &mut StagingDir, output_dir: &Path) -> Result<(), SurgeryError> { + if let Some(file) = find_foreign_safetensors_file(output_dir)? { + return Err(anyhow!( + "split-mtp: {} holds weight file {file}; --force does not delete foreign weight files", + output_dir.display() + ) + .into()); + } + for name in OUTPUT_FILES { + let destination = output_dir.join(name); + if std::fs::symlink_metadata(&destination) + .is_ok_and(|metadata| metadata.file_type().is_symlink()) + { + return Err(anyhow!( + "split-mtp: refusing symlinked output file {}", + destination.display() + ) + .into()); + } + if destination.is_dir() { + return Err(anyhow!( + "split-mtp: output path {} is a directory", + destination.display() + ) + .into()); + } + } + + let stale_index = output_dir.join(STALE_INDEX); + match std::fs::symlink_metadata(&stale_index) { + Ok(metadata) if metadata.is_dir() => { + return Err(anyhow!( + "split-mtp: stale index path {} is a directory", + stale_index.display() + ) + .into()); + } + Ok(_) => std::fs::remove_file(&stale_index)?, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => return Err(error.into()), + } + + for name in OUTPUT_FILES { + let source = staging.path().join(name); + if std::fs::symlink_metadata(&source).is_ok() { + std::fs::rename(&source, output_dir.join(name)).with_context(|| { + format!( + "split-mtp: committing staged {} to {}", + name, + output_dir.display() + ) + })?; + } + } + std::fs::remove_dir(staging.path())?; + staging.committed = true; + Ok(()) +} + +fn find_foreign_safetensors_file(dir: &Path) -> Result, SurgeryError> { + let mut found = Vec::new(); + for entry in std::fs::read_dir(dir)? { + let name = entry?.file_name().to_string_lossy().into_owned(); + if name != "model.safetensors" && name.ends_with(".safetensors") { + found.push(name); + } + } + found.sort(); + Ok(found.into_iter().next()) +} diff --git a/src/lib/mlxcel-surgery/src/ops/split_mtp_tests.rs b/src/lib/mlxcel-surgery/src/ops/split_mtp_tests.rs index 8c85ad5b2..1b8d531df 100644 --- a/src/lib/mlxcel-surgery/src/ops/split_mtp_tests.rs +++ b/src/lib/mlxcel-surgery/src/ops/split_mtp_tests.rs @@ -20,6 +20,7 @@ use std::path::Path; use mlxcel_core::weights::WeightMap; use serde_json::{Value, json}; +use super::output::{StagingDir, commit_staging_dir}; use super::{DRAFTER_MODEL_TYPE, SplitMtpOptions, split_mtp, split_mtp_dir, write_safetensors}; const VOCAB: i32 = 16; @@ -342,6 +343,7 @@ fn split_quantizes_projections_but_not_the_router_or_norms() { block_size: Some(3), q_bits: Some(4), q_group_size: 64, + force: false, }; let result = split_mtp(synthetic_weights(), &source_config(), &opts).expect("split"); let w = &result.weights; @@ -449,6 +451,7 @@ fn split_dir_writes_a_loadable_drafter_directory() { block_size: None, q_bits: Some(4), q_group_size: 64, + force: false, }; let report = split_mtp_dir(&src, &out, &opts).expect("split dir"); assert_eq!(report.source_layer, LAYERS); @@ -490,6 +493,254 @@ fn split_dir_writes_a_loadable_drafter_directory() { assert!(out.join("chat_template.jinja").is_file()); } +fn assert_no_staging_directory(parent: &Path, output_name: &str) { + let prefix = format!(".{output_name}.split-mtp-"); + let staged: Vec = std::fs::read_dir(parent) + .expect("read parent") + .flatten() + .map(|entry| entry.file_name().to_string_lossy().into_owned()) + .filter(|name| name.starts_with(&prefix)) + .collect(); + assert!(staged.is_empty(), "staging directories remain: {staged:?}"); +} + +#[test] +fn empty_existing_output_is_replaced_by_the_staged_directory() { + let tmp = tempfile::tempdir().unwrap(); + let out = tmp.path().join("drafter"); + std::fs::create_dir(&out).unwrap(); + let staging = StagingDir::create(&out).expect("staging"); + std::fs::write(staging.path().join("config.json"), b"new").unwrap(); + + commit_staging_dir(staging, &out, false).expect("commit"); + + assert_eq!(std::fs::read(out.join("config.json")).unwrap(), b"new"); + assert_no_staging_directory(tmp.path(), "drafter"); +} + +#[cfg(unix)] +#[test] +fn staging_directory_is_private() { + use std::os::unix::fs::PermissionsExt; + + let tmp = tempfile::tempdir().unwrap(); + let out = tmp.path().join("drafter"); + let staging = StagingDir::create(&out).expect("staging"); + + let mode = std::fs::metadata(staging.path()) + .expect("metadata") + .permissions() + .mode() + & 0o777; + assert_eq!(mode, 0o700); +} + +#[test] +fn dot_and_dotdot_outputs_stage_beside_the_resolved_directory() { + let tmp = tempfile::tempdir().unwrap(); + let dot_target = tmp.path().join("dot-target"); + let dotdot_child = tmp.path().join("dotdot-target").join("child"); + std::fs::create_dir(&dot_target).unwrap(); + std::fs::create_dir_all(&dotdot_child).unwrap(); + + for output in [dot_target.join("."), dotdot_child.join("..")] { + let staging = StagingDir::create(&output).expect("staging"); + assert_eq!(staging.path().parent(), Some(tmp.path())); + std::fs::write(staging.path().join("config.json"), b"new").unwrap(); + commit_staging_dir(staging, &output, true).expect("commit"); + assert_eq!(std::fs::read(output.join("config.json")).unwrap(), b"new"); + } +} + +#[cfg(unix)] +#[test] +fn non_utf8_output_name_is_supported() { + use std::os::unix::ffi::OsStringExt; + + let tmp = tempfile::tempdir().unwrap(); + let output = tmp.path().join(std::ffi::OsString::from_vec(vec![ + b'd', b'r', 0xff, b'f', b't', + ])); + let staging = StagingDir::create(&output).expect("staging"); + std::fs::write(staging.path().join("config.json"), b"new").unwrap(); + + commit_staging_dir(staging, &output, false).expect("commit"); + + assert_eq!(std::fs::read(output.join("config.json")).unwrap(), b"new"); +} + +#[cfg(unix)] +#[test] +fn commit_refuses_an_ancestor_symlink_retarget() { + use std::os::unix::fs::symlink; + + let tmp = tempfile::tempdir().unwrap(); + let first = tmp.path().join("first"); + let second = tmp.path().join("second"); + let route = tmp.path().join("route"); + std::fs::create_dir(&first).unwrap(); + std::fs::create_dir(&second).unwrap(); + symlink(&first, &route).unwrap(); + let output = route.join("drafter"); + let staging = StagingDir::create(&output).expect("staging"); + std::fs::write(staging.path().join("config.json"), b"new").unwrap(); + std::fs::remove_file(&route).unwrap(); + symlink(&second, &route).unwrap(); + + let err = commit_staging_dir(staging, &output, false).expect_err("must refuse drift"); + + assert!(err.to_string().contains("changed while"), "{err}"); + assert!(!first.join("drafter").exists()); + assert!(!second.join("drafter").exists()); + assert_no_staging_directory(&first, "drafter"); +} + +#[test] +fn concurrent_non_empty_output_wins_without_interleaving() { + let tmp = tempfile::tempdir().unwrap(); + let out = tmp.path().join("drafter"); + let staging = StagingDir::create(&out).expect("staging"); + std::fs::write(staging.path().join("config.json"), b"loser").unwrap(); + std::fs::create_dir(&out).unwrap(); + std::fs::write(out.join("config.json"), b"winner").unwrap(); + std::fs::write(out.join("model.safetensors"), b"winner weights").unwrap(); + + let err = commit_staging_dir(staging, &out, false).expect_err("must refuse"); + + assert!(err.to_string().contains("config.json"), "{err}"); + assert_eq!(std::fs::read(out.join("config.json")).unwrap(), b"winner"); + assert_eq!( + std::fs::read(out.join("model.safetensors")).unwrap(), + b"winner weights" + ); + assert_no_staging_directory(tmp.path(), "drafter"); +} + +#[test] +fn split_dir_refuses_an_unrelated_non_empty_output_without_force() { + let tmp = tempfile::tempdir().unwrap(); + let src = tmp.path().join("raw"); + let out = tmp.path().join("drafter"); + std::fs::create_dir(&src).unwrap(); + std::fs::create_dir(&out).unwrap(); + write_synthetic_checkpoint(&src); + std::fs::write(out.join("README.md"), b"keep").unwrap(); + + let err = split_mtp_dir(&src, &out, &SplitMtpOptions::default()).expect_err("must refuse"); + + assert!(err.to_string().contains("README.md"), "{err}"); + assert_eq!(std::fs::read(out.join("README.md")).unwrap(), b"keep"); + assert!(!out.join("model.safetensors").exists()); + assert_no_staging_directory(tmp.path(), "drafter"); +} + +#[cfg(unix)] +#[test] +fn split_dir_refuses_present_and_dangling_output_symlinks_for_both_force_modes() { + use std::os::unix::fs::symlink; + + let tmp = tempfile::tempdir().unwrap(); + let src = tmp.path().join("raw"); + std::fs::create_dir(&src).unwrap(); + write_synthetic_checkpoint(&src); + + for force in [false, true] { + let suffix = if force { "force" } else { "plain" }; + let tokenizer_out = tmp.path().join(format!("tokenizer-{suffix}")); + let tokenizer_target = tmp.path().join(format!("tokenizer-target-{suffix}")); + std::fs::create_dir(&tokenizer_out).unwrap(); + std::fs::write(&tokenizer_target, b"untouched").unwrap(); + symlink(&tokenizer_target, tokenizer_out.join("tokenizer.json")).unwrap(); + let opts = SplitMtpOptions { + force, + ..SplitMtpOptions::default() + }; + let err = split_mtp_dir(&src, &tokenizer_out, &opts).expect_err("must refuse"); + assert!(err.to_string().contains("tokenizer.json"), "{err}"); + assert_eq!(std::fs::read(&tokenizer_target).unwrap(), b"untouched"); + assert_no_staging_directory(tmp.path(), &format!("tokenizer-{suffix}")); + + let config_out = tmp.path().join(format!("config-{suffix}")); + let config_target = tmp.path().join(format!("config-target-{suffix}")); + std::fs::create_dir(&config_out).unwrap(); + symlink(&config_target, config_out.join("config.json")).unwrap(); + let err = split_mtp_dir(&src, &config_out, &opts).expect_err("must refuse"); + assert!(err.to_string().contains("config.json"), "{err}"); + assert!(!config_target.exists(), "dangling target must stay absent"); + assert_no_staging_directory(tmp.path(), &format!("config-{suffix}")); + } +} + +#[test] +fn force_commit_replaces_drafter_files_and_removes_the_stale_index() { + let tmp = tempfile::tempdir().unwrap(); + let src = tmp.path().join("raw"); + let out = tmp.path().join("drafter"); + std::fs::create_dir(&src).unwrap(); + std::fs::create_dir(&out).unwrap(); + write_synthetic_checkpoint(&src); + std::fs::write(out.join("model.safetensors"), b"old").unwrap(); + std::fs::write(out.join("config.json"), b"old").unwrap(); + std::fs::write(out.join("model.safetensors.index.json"), b"stale").unwrap(); + let opts = SplitMtpOptions { + force: true, + ..SplitMtpOptions::default() + }; + + split_mtp_dir(&src, &out, &opts).expect("force replacement"); + + assert!(!out.join("model.safetensors.index.json").exists()); + let config: Value = + serde_json::from_str(&std::fs::read_to_string(out.join("config.json")).unwrap()).unwrap(); + assert_eq!(config["model_type"], DRAFTER_MODEL_TYPE); + assert_no_staging_directory(tmp.path(), "drafter"); +} + +#[test] +fn safetensors_writer_refuses_to_replace_an_existing_file() { + let tmp = tempfile::tempdir().unwrap(); + let path = tmp.path().join("model.safetensors"); + std::fs::write(&path, b"keep").unwrap(); + + let err = write_safetensors(synthetic_weights(), &path).expect_err("must refuse"); + + assert!(err.to_string().contains("creating"), "{err}"); + assert_eq!(std::fs::read(path).unwrap(), b"keep"); +} + +#[test] +fn exclusive_safetensors_writer_matches_the_crate_serializer_byte_for_byte() { + let tmp = tempfile::tempdir().unwrap(); + let exclusive_path = tmp.path().join("exclusive.safetensors"); + let reference_path = tmp.path().join("reference.safetensors"); + let mut weights = WeightMap::new(); + weights.insert("z-f32".to_string(), ramp(&[2, 3], 1.0)); + weights.insert("a-f32".to_string(), ramp(&[1, 4], 2.0)); + let bf16_source = ramp(&[2, 2], 3.0); + weights.insert( + "m-bf16".to_string(), + mlxcel_core::astype(&bf16_source, mlxcel_core::dtype::BFLOAT16), + ); + let u32_source = mlxcel_core::from_slice_i32(&[1, 2, 3], &[3]); + weights.insert( + "b-u32".to_string(), + mlxcel_core::astype(&u32_source, mlxcel_core::dtype::UINT32), + ); + + write_safetensors(weights, &exclusive_path).expect("exclusive write"); + let exclusive = std::fs::read(&exclusive_path).expect("read exclusive"); + let parsed = safetensors::SafeTensors::deserialize(&exclusive).expect("parse exclusive"); + let mut metadata = HashMap::new(); + metadata.insert("format".to_string(), "mlx".to_string()); + safetensors::serialize_to_file(parsed.tensors(), Some(metadata), &reference_path) + .expect("reference write"); + + assert_eq!( + exclusive, + std::fs::read(reference_path).expect("read reference") + ); +} + #[test] fn split_dir_refuses_an_output_that_aliases_the_source() { // `--output ` passes the CLI overwrite guard on a sharded @@ -522,6 +773,40 @@ fn split_dir_refuses_an_output_that_aliases_the_source() { assert!(nested.join("model.safetensors").is_file()); } +#[test] +fn split_dir_refuses_an_unresolved_dotdot_alias_without_touching_the_source() { + let tmp = tempfile::tempdir().unwrap(); + let src = tmp.path().join("raw"); + std::fs::create_dir(&src).unwrap(); + write_synthetic_checkpoint(&src); + let config_before = std::fs::read(src.join("config.json")).unwrap(); + let shard = src.join("model-00001-of-00003.safetensors"); + let weights_before = std::fs::read(&shard).unwrap(); + let mut entries_before: Vec<_> = std::fs::read_dir(&src) + .unwrap() + .map(|entry| entry.unwrap().file_name()) + .collect(); + entries_before.sort(); + let output = src.join("new").join(".."); + + let err = split_mtp_dir(&src, &output, &SplitMtpOptions::default()).expect_err("must refuse"); + + assert!(err.to_string().contains("source checkpoint"), "{err}"); + assert_eq!( + std::fs::read(src.join("config.json")).unwrap(), + config_before + ); + assert_eq!(std::fs::read(shard).unwrap(), weights_before); + let mut entries_after: Vec<_> = std::fs::read_dir(&src) + .unwrap() + .map(|entry| entry.unwrap().file_name()) + .collect(); + entries_after.sort(); + assert_eq!(entries_after, entries_before); + assert!(!src.join("new").exists()); + assert_no_staging_directory(tmp.path(), "raw"); +} + #[test] fn split_dir_refuses_an_index_without_the_nextn_layer() { let tmp = tempfile::tempdir().unwrap();