Make quantized layer-package builds resumable

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.

crates/skippy-model-package/src/main.rs

@@ -71,13 +71,17 @@ enum Command {
         #[arg(long)]
         after_artifact_command: Option<PathBuf>,
         #[arg(long)]
+        transform_artifact_command: Option<PathBuf>,
+        #[arg(long)]
         model_id: Option<String>,
         #[arg(long)]
         source_repo: Option<String>,
         #[arg(long)]
         source_revision: Option<String>,
         #[arg(long)]
         source_file: Option<String>,
+        #[arg(long)]
+        resume_existing_artifacts: bool,
     },
     Validate {
         full: PathBuf,
@@ -405,23 +409,29 @@ 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,
             projectors,
             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<PathBuf>,
     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<PackageArtifact> {
     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<PackageArtifact> {
+    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<u32>) -> TensorInfo {
         TensorInfo {
             name: name.to_string(),

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:

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\""));
+    }
 }