diff --git a/crates/skippy-model-package/README.md b/crates/skippy-model-package/README.md index 5590dbcdec..fc7580dd8a 100644 --- a/crates/skippy-model-package/README.md +++ b/crates/skippy-model-package/README.md @@ -88,6 +88,13 @@ skippy-model-package write-package ./model.gguf \ This keeps canonical package identity tied to real model coordinates rather than inferred from arbitrary filesystem paths. +`--transform-artifact-command` runs an in-place transformation before package +metadata is measured, so the manifest describes the transformed artifact. +`--after-artifact-command` remains suitable for upload hooks that may remove +the artifact after its metadata is captured. Pass `--resume-existing-artifacts` +to reuse existing artifacts; transform and upload hooks are still run for each +resumed artifact. + `validate-package` checks the source-model checksum, manifest artifact checksums and sizes, declared tensor counts/bytes, layer coverage, duplicate layers, and exact owned tensor coverage against the source model. diff --git a/crates/skippy-model-package/src/main.rs b/crates/skippy-model-package/src/main.rs index d2fb273553..60985d0bf6 100644 --- a/crates/skippy-model-package/src/main.rs +++ b/crates/skippy-model-package/src/main.rs @@ -71,6 +71,8 @@ enum Command { #[arg(long)] after_artifact_command: Option, #[arg(long)] + transform_artifact_command: Option, + #[arg(long)] model_id: Option, #[arg(long)] source_repo: Option, @@ -78,6 +80,8 @@ enum Command { source_revision: Option, #[arg(long)] source_file: Option, + #[arg(long)] + resume_existing_artifacts: bool, }, Validate { full: PathBuf, @@ -405,10 +409,12 @@ fn main() -> Result<()> { out_dir, projectors, after_artifact_command, + transform_artifact_command, model_id, source_repo, source_revision, source_file, + resume_existing_artifacts, } => write_package( model, out_dir, @@ -416,12 +422,16 @@ fn main() -> Result<()> { ArtifactHook { command: after_artifact_command, }, + ArtifactHook { + command: transform_artifact_command, + }, ExplicitSourceIdentity { model_id, source_repo, source_revision, source_file, }, + resume_existing_artifacts, ), Command::Validate { full, slices } => validate(full, slices), Command::ValidatePackage { full, package } => validate_package(full, package), @@ -517,7 +527,9 @@ fn write_package( out_dir: PathBuf, projectors: Vec, artifact_hook: ArtifactHook, + artifact_transform: ArtifactHook, explicit: ExplicitSourceIdentity, + resume_existing_artifacts: bool, ) -> Result<()> { let input = resolve_package_input(model, explicit)?; fs::create_dir_all(&out_dir) @@ -552,6 +564,8 @@ fn write_package( }, &out_dir, &artifact_hook, + &artifact_transform, + resume_existing_artifacts, )?; progress.finish_step(&artifact_progress_detail(&metadata))?; progress.start_step("shared/embeddings.gguf")?; @@ -568,6 +582,8 @@ fn write_package( }, &out_dir, &artifact_hook, + &artifact_transform, + resume_existing_artifacts, )?; progress.finish_step(&artifact_progress_detail(&embeddings))?; progress.start_step("shared/output.gguf")?; @@ -584,6 +600,8 @@ fn write_package( }, &out_dir, &artifact_hook, + &artifact_transform, + resume_existing_artifacts, )?; progress.finish_step(&artifact_progress_detail(&output))?; @@ -604,6 +622,8 @@ fn write_package( }, &out_dir, &artifact_hook, + &artifact_transform, + resume_existing_artifacts, )?; progress.finish_step(&artifact_progress_detail(&artifact))?; layers.push(PackageLayer { @@ -1178,6 +1198,8 @@ fn write_package_artifact( spec: PackageArtifactSpec, out_dir: &Path, artifact_hook: &ArtifactHook, + artifact_transform: &ArtifactHook, + resume_existing_artifacts: bool, ) -> Result { let stage = stage_plan_from_tensors( spec.stage_index as usize, @@ -1188,28 +1210,35 @@ fn write_package_artifact( tensors, ); let path = out_dir.join(&spec.relative_path); - write_stage_artifact(source, &stage, &path)?; + if !should_resume_package_artifact(&path, resume_existing_artifacts) { + write_stage_artifact(source, &stage, &path)?; + } let relative_path = spec.relative_path.display().to_string(); - // Read all artifact metadata from disk before running the hook: the upload - // hook may move or delete the local file (see split-model-job.sh, which - // uploads then unlinks each artifact to stay under the ephemeral storage - // limit). Reopening the file after the hook would fail once it is gone. - let artifact_info = ModelInfo::open(&path) + run_artifact_hook(artifact_transform, &path, &relative_path)?; + let artifact = read_package_artifact(&path, &spec.relative_path)?; + run_artifact_hook(artifact_hook, &path, &artifact.path)?; + Ok(artifact) +} + +fn should_resume_package_artifact(path: &Path, resume_existing_artifacts: bool) -> bool { + resume_existing_artifacts && path.is_file() +} + +fn read_package_artifact(path: &Path, relative_path: &Path) -> Result { + let artifact_info = ModelInfo::open(path) .with_context(|| format!("open package artifact {}", path.display()))?; let artifact_tensors = artifact_info .tensors() .with_context(|| format!("read package artifact tensors {}", path.display()))?; - let metadata = fs::metadata(&path) - .with_context(|| format!("read artifact metadata {}", path.display()))?; - let artifact = PackageArtifact { - path: relative_path.clone(), + let metadata = + fs::metadata(path).with_context(|| format!("read artifact metadata {}", path.display()))?; + Ok(PackageArtifact { + path: relative_path.display().to_string(), tensor_count: artifact_tensors.len(), tensor_bytes: artifact_tensors.iter().map(|tensor| tensor.byte_size).sum(), artifact_bytes: metadata.len(), - sha256: file_sha256(&path)?, - }; - run_artifact_hook(artifact_hook, &path, &relative_path)?; - Ok(artifact) + sha256: file_sha256(path)?, + }) } fn copy_projector_artifact( @@ -1955,6 +1984,7 @@ mod tests { ArtifactHook, ExplicitSourceIdentity, activation_width, local_artifact_files, model_distribution_id, native_mtp_layer_indices, package_generation, resolve_gguf_shard_paths, resolve_local_package_input, run_artifact_hook, + should_resume_package_artifact, }; use skippy_ffi::TensorRole; use skippy_runtime::TensorInfo; @@ -2245,6 +2275,22 @@ mod tests { std::fs::remove_dir_all(dir).unwrap(); } + #[test] + fn resumes_only_existing_artifacts_when_requested() { + let dir = unique_test_dir("resume-artifact"); + std::fs::create_dir_all(&dir).unwrap(); + let artifact = dir.join("layer-000.gguf"); + std::fs::write(&artifact, b"existing").unwrap(); + + assert!(should_resume_package_artifact(&artifact, true)); + assert!(!should_resume_package_artifact(&artifact, false)); + assert!(!should_resume_package_artifact( + &dir.join("missing.gguf"), + true + )); + std::fs::remove_dir_all(dir).unwrap(); + } + fn tensor(name: &str, layer_index: Option) -> TensorInfo { TensorInfo { name: name.to_string(), diff --git a/crates/skippy-quantize/README.md b/crates/skippy-quantize/README.md index 4451c293a3..c1e452196b 100644 --- a/crates/skippy-quantize/README.md +++ b/crates/skippy-quantize/README.md @@ -471,6 +471,13 @@ pass `--keep-quant` to retain it. It does not pass `--max-memory` to quantization because the unpatched llama API quant backend does not expose a memory-budget knob. +Pass `--resume-package` to reuse package artifacts already present in +`--package-dir`. Each artifact uses its own quantization manifest, so an +interrupted package build resumes at the first missing artifact. The package +transform hook keeps its completed output in quant scratch and atomically +links it into the package, avoiding a second model-sized copy while preserving +a resumable result until package preflight succeeds. + ## Validation Useful checks: diff --git a/crates/skippy-quantize/src/main.rs b/crates/skippy-quantize/src/main.rs index 90abfc49e2..a64cd93fad 100644 --- a/crates/skippy-quantize/src/main.rs +++ b/crates/skippy-quantize/src/main.rs @@ -227,6 +227,8 @@ struct QuantizeLayerPackageArgs { keep_quant: bool, #[arg(long)] replace_package: bool, + #[arg(long, conflicts_with = "replace_package")] + resume_package: bool, #[arg(long)] json: bool, } @@ -694,7 +696,7 @@ fn quantize_layer_package(args: QuantizeLayerPackageArgs) -> Result<()> { "missing skippy-model-package binary {}; build it with `cargo build --release --locked -p skippy-model-package` or pass --skippy-model-package-bin", args.skippy_model_package_bin.display() ); - if args.package_dir.exists() { + if args.package_dir.exists() && !args.resume_package { ensure!( args.replace_package, "package dir already exists: {}; pass --replace-package to overwrite it", @@ -750,15 +752,26 @@ fn write_layer_package_quant_hook( let hook_spool_dir = args.init.target.join("hook-spool"); let hook_record_dir = args.init.target.join("hook-records"); let hook_status_file = args.init.target.join("hook-status.json"); + let hook_manifest_dir = args.init.target.join("hook-manifests"); + let hook_result_dir = args.init.target.join("hook-results"); let mut script = String::new(); script.push_str("#!/usr/bin/env bash\nset -euo pipefail\n"); script.push_str("case \"${SKIPPY_PACKAGE_ARTIFACT_RELATIVE_PATH:-}\" in\n"); script.push_str(" shared/metadata.gguf) exit 0 ;;\n"); script.push_str("esac\n"); script.push_str("artifact=\"${SKIPPY_PACKAGE_ARTIFACT_PATH:?}\"\n"); - script.push_str("tmp=\"${artifact}.quant-tmp.gguf\"\n"); + script.push_str("artifact_rel=\"${SKIPPY_PACKAGE_ARTIFACT_RELATIVE_PATH:-artifact}\"\n"); + script.push_str("artifact_safe=\"${artifact_rel//\\//_}\"\n"); + script.push_str(&format!( + "manifest_dir={}\n", + shell_quote(&hook_manifest_dir) + )); + script.push_str(&format!("result_dir={}\n", shell_quote(&hook_result_dir))); + script.push_str("mkdir -p \"$manifest_dir\" \"$result_dir\"\n"); + script.push_str("manifest=\"$manifest_dir/${artifact_safe}.skippy-quantize.json\"\n"); + script.push_str("result=\"$result_dir/${artifact_safe}.gguf\"\n"); script.push_str("single_source=\"${artifact}.quant-src\"\n"); - script.push_str("rm -rf \"$single_source\" \"$tmp\"\n"); + script.push_str("rm -rf \"$single_source\"\n"); script.push_str("mkdir -p \"$single_source\"\n"); script.push_str("ln -s \"$artifact\" \"$single_source/model.gguf\"\n"); script.push_str(&format!( @@ -772,6 +785,7 @@ fn write_layer_package_quant_hook( runner.json_event_interval_seconds, runner.json_event_window, )); + append_layer_package_hook_quant_flags(&mut script); if let Some(watchdog_seconds) = runner.watchdog_seconds { script.push_str(&format!(" --watchdog-seconds {watchdog_seconds}")); } @@ -795,14 +809,16 @@ fn write_layer_package_quant_hook( shell_quote(library) )); } - script.push_str(" \"$single_source/model.gguf\" \"$tmp\" "); + script.push_str(" \"$single_source/model.gguf\" \"$result\" "); script.push_str(&shell_quote(args.init.quant.output_name())); script.push('\n'); script.push_str("rm -rf \"$single_source\"\n"); - script.push_str("if [[ ! -f \"$tmp\" && -f \"${tmp%.gguf}-00001-of-00001.gguf\" ]]; then\n"); - script.push_str(" tmp=\"${tmp%.gguf}-00001-of-00001.gguf\"\n"); + script.push_str( + "if [[ ! -f \"$result\" && -f \"${result%.gguf}-00001-of-00001.gguf\" ]]; then\n", + ); + script.push_str(" result=\"${result%.gguf}-00001-of-00001.gguf\"\n"); script.push_str("fi\n"); - script.push_str("mv \"$tmp\" \"$artifact\"\n"); + append_layer_package_result_install(&mut script); fs::write(&hook, script).with_context(|| format!("write hook {}", hook.display()))?; make_executable(&hook)?; Ok(hook) @@ -814,6 +830,17 @@ fn append_optional_tensor_type_file(script: &mut String, tensor_type_file: Optio } } +fn append_layer_package_hook_quant_flags(script: &mut String) { + script.push_str(" --manifest \"$manifest\" --no-stage-source"); +} + +fn append_layer_package_result_install(script: &mut String) { + script.push_str("ready=\"${artifact}.quant-ready\"\n"); + script.push_str("rm -f \"$ready\"\n"); + script.push_str("ln \"$result\" \"$ready\" 2>/dev/null || cp \"$result\" \"$ready\"\n"); + script.push_str("mv -f \"$ready\" \"$artifact\"\n"); +} + fn append_override_kv_args(script: &mut String, overrides: &[String]) { for override_kv in overrides { script.push_str(&format!(" --override-kv {}", shell_quote_str(override_kv))); @@ -833,12 +860,13 @@ fn run_skippy_model_package_write( .unwrap_or("model.gguf"); format!("{}/{}", args.init.source_prefix, file_name) }); - let status = ProcessCommand::new(&args.skippy_model_package_bin) + let mut command = ProcessCommand::new(&args.skippy_model_package_bin); + command .arg("write-package") .arg(first_source_shard) .arg("--out-dir") .arg(&args.package_dir) - .arg("--after-artifact-command") + .arg("--transform-artifact-command") .arg(hook) .arg("--model-id") .arg(&args.package_model_id) @@ -847,14 +875,16 @@ fn run_skippy_model_package_write( .arg("--source-revision") .arg(&args.package_source_revision) .arg("--source-file") - .arg(source_file) - .status() - .with_context(|| { - format!( - "run {} write-package", - args.skippy_model_package_bin.display() - ) - })?; + .arg(source_file); + if args.resume_package { + command.arg("--resume-existing-artifacts"); + } + let status = command.status().with_context(|| { + format!( + "run {} write-package", + args.skippy_model_package_bin.display() + ) + })?; ensure!( status.success(), "{} write-package failed with status {status}", @@ -1496,7 +1526,10 @@ fn print_dry_run_complete(json: bool, kind: &str) -> Result<()> { mod tests { use std::path::Path; - use super::{append_optional_tensor_type_file, append_override_kv_args}; + use super::{ + append_layer_package_hook_quant_flags, append_layer_package_result_install, + append_optional_tensor_type_file, append_override_kv_args, + }; #[test] fn layer_package_hook_forwards_tensor_type_file() { @@ -1522,4 +1555,24 @@ mod tests { assert!(script.contains("--override-kv 'glm-dsa.attention.indexer.head_count=int:32'")); } + + #[test] + fn layer_package_hook_skips_redundant_source_staging() { + let mut script = String::new(); + + append_layer_package_hook_quant_flags(&mut script); + + assert!(script.contains("--manifest \"$manifest\"")); + assert!(script.contains("--no-stage-source")); + } + + #[test] + fn layer_package_hook_installs_persistent_result_atomically() { + let mut script = String::new(); + + append_layer_package_result_install(&mut script); + + assert!(script.contains("ln \"$result\" \"$ready\"")); + assert!(script.contains("mv -f \"$ready\" \"$artifact\"")); + } }