model_package/
script.rs

1use anyhow::{Context, Result};
2use hf_hub::{
3    HFClient,
4    buckets::{BucketDownload, BucketTreeEntry, BucketUpload},
5};
6use sha2::{Digest, Sha256};
7
8/// The embedded canonical job script.
9pub const EMBEDDED_SCRIPT: &str = include_str!("scripts/split-model-job.sh");
10
11// Generated by build.rs — provides EMBEDDED_SCRIPT_SHA256 and EMBEDDED_SCRIPT_SIZE.
12include!(concat!(env!("OUT_DIR"), "/script_hash.rs"));
13
14/// Result of comparing the bucket script against the embedded version.
15#[derive(Debug)]
16pub struct ScriptFreshness {
17    /// Size of the script in the bucket (bytes).
18    pub bucket_size: u64,
19    /// Size of the embedded script (bytes).
20    pub expected_size: u64,
21    /// Whether the bucket script matches the embedded version.
22    pub is_current: bool,
23    /// Human-readable reason when the script is not current.
24    pub mismatch_reason: Option<String>,
25}
26
27/// Compare the bucket script's size against the embedded version.
28///
29/// If sizes match, downloads the ~10 KB file and verifies the SHA-256 hash.
30pub async fn check_bucket_script(client: &HFClient) -> Result<ScriptFreshness> {
31    let bucket = client.bucket("meshllm", "layer-split-output");
32    let paths = vec!["split-model-job.sh".to_string()];
33
34    let entries = bucket
35        .get_paths_info()
36        .paths(paths)
37        .send()
38        .await
39        .context("check bucket script metadata")?;
40
41    let bucket_size = entries
42        .into_iter()
43        .find_map(|entry| match entry {
44            BucketTreeEntry::File { size, .. } => Some(size),
45            _ => None,
46        })
47        .context("split-model-job.sh not found in meshllm/layer-split-output bucket")?;
48
49    if bucket_size != EMBEDDED_SCRIPT_SIZE {
50        return Ok(ScriptFreshness {
51            bucket_size,
52            expected_size: EMBEDDED_SCRIPT_SIZE,
53            is_current: false,
54            mismatch_reason: Some(format!(
55                "size differs: bucket has {bucket_size} bytes, embedded script has {expected_size} bytes",
56                expected_size = EMBEDDED_SCRIPT_SIZE
57            )),
58        });
59    }
60
61    // Sizes match — download to a unique temp file and verify SHA-256.
62    let tmp_dir = tempfile::tempdir().context("create bucket script check temp dir")?;
63    let tmp_path = tmp_dir.path().join("split-model-job.sh");
64
65    bucket
66        .download_files()
67        .files(vec![BucketDownload::new(
68            "split-model-job.sh",
69            tmp_path.clone(),
70        )])
71        .send()
72        .await
73        .context("download bucket script for hash verification")?;
74
75    let bytes = std::fs::read(&tmp_path).context("read downloaded bucket script")?;
76    let hash = Sha256::digest(&bytes);
77    let hex: String = hash.iter().map(|b| format!("{b:02x}")).collect();
78
79    let is_current = hex == EMBEDDED_SCRIPT_SHA256;
80
81    Ok(ScriptFreshness {
82        bucket_size,
83        expected_size: EMBEDDED_SCRIPT_SIZE,
84        is_current,
85        mismatch_reason: (!is_current).then(|| {
86            format!(
87                "sha256 differs: bucket has {hex}, embedded script has {expected}",
88                expected = EMBEDDED_SCRIPT_SHA256
89            )
90        }),
91    })
92}
93
94/// Upload the embedded script to the meshllm bucket.
95///
96/// Requires meshllm org write access.
97pub async fn update_bucket_script(client: &HFClient) -> Result<()> {
98    let bucket = client.bucket("meshllm", "layer-split-output");
99
100    // Write embedded script to a unique temp file for upload.
101    let tmp_dir = tempfile::tempdir().context("create bucket script upload temp dir")?;
102    let tmp_path = tmp_dir.path().join("split-model-job.sh");
103    std::fs::write(&tmp_path, EMBEDDED_SCRIPT)?;
104
105    bucket
106        .upload_files()
107        .files(vec![BucketUpload::new(
108            tmp_path.clone(),
109            "split-model-job.sh",
110        )])
111        .send()
112        .await
113        .context("upload script to meshllm/layer-split-output bucket")?;
114
115    Ok(())
116}
117
118#[cfg(test)]
119mod tests {
120    use super::EMBEDDED_SCRIPT;
121
122    #[test]
123    fn embedded_script_prints_mesh_llm_repo_shorthand() {
124        assert!(EMBEDDED_SCRIPT.contains(r#"mesh-llm serve --model "{target_repo}" --split"#));
125        assert!(EMBEDDED_SCRIPT.contains("mesh-llm serve --model ${TARGET_REPO} --split"));
126        assert!(!EMBEDDED_SCRIPT.contains(r#"mesh-llm serve --model "hf://${TARGET_REPO}""#));
127        assert!(!EMBEDDED_SCRIPT.contains("mesh-llm serve --model hf://${TARGET_REPO}"));
128    }
129
130    #[test]
131    fn embedded_script_writes_rich_model_card() {
132        assert!(EMBEDDED_SCRIPT.contains("SOURCE_PIPELINE_TAG"));
133        assert!(EMBEDDED_SCRIPT.contains("PACKAGE_EXPERIMENTAL"));
134        assert!(EMBEDDED_SCRIPT.contains("pipeline_tag: {yaml_quote(source_pipeline_tag)}"));
135        assert!(EMBEDDED_SCRIPT.contains("- experimental"));
136        assert!(EMBEDDED_SCRIPT.contains("Experimental package:"));
137        assert!(EMBEDDED_SCRIPT.contains("resolve_upstream_license"));
138        assert!(EMBEDDED_SCRIPT.contains("license_frontmatter"));
139        assert!(EMBEDDED_SCRIPT.contains("could not resolve upstream license metadata"));
140        assert!(EMBEDDED_SCRIPT.contains("- openai-compatible"));
141        assert!(EMBEDDED_SCRIPT.contains("## Model Overview"));
142        assert!(EMBEDDED_SCRIPT.contains("## Highlights"));
143        assert!(EMBEDDED_SCRIPT.contains("## Recommended Use"));
144        assert!(EMBEDDED_SCRIPT.contains("## Quickstart"));
145        assert!(EMBEDDED_SCRIPT.contains("## Package Variant"));
146        assert!(EMBEDDED_SCRIPT.contains("## What Is Included"));
147        assert!(EMBEDDED_SCRIPT.contains("## Validation"));
148        assert!(
149            EMBEDDED_SCRIPT.contains("https://meshllm.cloud/assets/images/jelly-logo-wordmark.png")
150        );
151        assert!(EMBEDDED_SCRIPT.contains("style=for-the-badge"));
152        assert!(EMBEDDED_SCRIPT.contains("https://www.meshllm.cloud"));
153        assert!(EMBEDDED_SCRIPT.contains("https://discord.gg/rs6fmc63eN"));
154        assert!(EMBEDDED_SCRIPT.contains("Distributed GGUF inference package for Mesh LLM"));
155        assert!(EMBEDDED_SCRIPT.contains(
156            "For upstream architecture details, chat template guidance, sampling recommendations"
157        ));
158        assert!(EMBEDDED_SCRIPT.contains("Package manifest SHA-256"));
159        assert!(EMBEDDED_SCRIPT.contains("uploaded to this repository"));
160    }
161
162    #[test]
163    fn embedded_script_streams_package_artifacts_to_hub() {
164        assert!(EMBEDDED_SCRIPT.contains("SOURCE_QUANT"));
165        assert!(EMBEDDED_SCRIPT.contains("SOURCE_TOTAL_BYTES"));
166        assert!(EMBEDDED_SCRIPT.contains("Estimated fallback /bucket cache needed"));
167        assert!(EMBEDDED_SCRIPT.contains("estimate_bucket_workspace_bytes"));
168        assert!(EMBEDDED_SCRIPT.contains(r#"HF_HUB_CACHE="${HF_HUB_CACHE:-${HF_HOME}/hub}""#));
169        assert!(EMBEDDED_SCRIPT.contains(r#"HF_XET_CACHE="${HF_XET_CACHE:-${HF_HOME}/xet}""#));
170        assert!(
171            EMBEDDED_SCRIPT.contains(r#"PACKAGE_DIR="${PACKAGE_DIR:-${LOCAL_WORK_DIR}/package}""#)
172        );
173        assert!(EMBEDDED_SCRIPT.contains(r#"JOB_TMP_DIR="${JOB_TMP_DIR:-${LOCAL_WORK_DIR}/tmp}""#));
174        assert!(EMBEDDED_SCRIPT.contains(r#"LOCAL_WORK_DIR="${LOCAL_WORK_DIR:-/tmp/"#));
175        assert!(
176            EMBEDDED_SCRIPT.contains(r#"BUILD_TMP_DIR="${BUILD_TMP_DIR:-${LOCAL_WORK_DIR}/tmp}""#)
177        );
178        assert!(EMBEDDED_SCRIPT.contains(r#"TMPDIR="$BUILD_TMP_DIR""#));
179        assert!(EMBEDDED_SCRIPT.contains(r#"TMPDIR="$JOB_TMP_DIR""#));
180        assert!(EMBEDDED_SCRIPT.contains("export TMPDIR TEMP TMP"));
181        assert!(
182            EMBEDDED_SCRIPT.contains(
183                r#"CARGO_TARGET_DIR="${CARGO_TARGET_DIR:-${LOCAL_WORK_DIR}/cargo-target}""#
184            )
185        );
186        assert!(
187            EMBEDDED_SCRIPT.contains(
188                r#"rm -rf "$BUILD_DIR" "$CARGO_TARGET_DIR" "$CARGO_HOME" "$RUSTUP_HOME""#
189            )
190        );
191        assert!(EMBEDDED_SCRIPT.contains(r#"SOURCE_REF="${SOURCE_REPO}:${SOURCE_QUANT}""#));
192        assert!(
193            EMBEDDED_SCRIPT
194                .contains(r#"SOURCE_REF="${SOURCE_REPO}@${SOURCE_REVISION}:${SOURCE_QUANT}""#)
195        );
196        assert!(EMBEDDED_SCRIPT.contains("log_storage_snapshot"));
197        assert!(EMBEDDED_SCRIPT.contains("start_heartbeat"));
198        assert!(EMBEDDED_SCRIPT.contains("Starting write-package"));
199        assert!(EMBEDDED_SCRIPT.contains("upload-package-artifact.py"));
200        assert!(EMBEDDED_SCRIPT.contains("SKIPPY_PACKAGE_ARTIFACT_PATH"));
201        assert!(EMBEDDED_SCRIPT.contains("SKIPPY_PACKAGE_ARTIFACT_RELATIVE_PATH"));
202        assert!(EMBEDDED_SCRIPT.contains(r#"--after-artifact-command "$ARTIFACT_UPLOAD_HOOK""#));
203        assert!(EMBEDDED_SCRIPT.contains("Uploaded and removed"));
204        assert!(!EMBEDDED_SCRIPT.contains("api.upload_folder"));
205        assert!(EMBEDDED_SCRIPT.contains(r#"MOUNTED_SOURCE_PATH="/source/${SOURCE_FILE}""#));
206        assert!(EMBEDDED_SCRIPT.contains(r#"WRITE_PACKAGE_INPUT="$MOUNTED_SOURCE_PATH""#));
207        assert!(EMBEDDED_SCRIPT.contains(r#"--source-file "$SOURCE_FILE""#));
208        assert!(EMBEDDED_SCRIPT.contains("SOURCE_PROJECTOR_FILES"));
209        assert!(
210            EMBEDDED_SCRIPT
211                .contains(r#"WRITE_PACKAGE_PROJECTOR_ARGS+=(--projector "$PROJECTOR_PATH")"#)
212        );
213        assert!(EMBEDDED_SCRIPT.contains(r#""${WRITE_PACKAGE_PROJECTOR_ARGS[@]}""#));
214        assert!(EMBEDDED_SCRIPT.contains(r#"time "$SLICER" write-package "$WRITE_PACKAGE_INPUT""#));
215        assert!(!EMBEDDED_SCRIPT.contains(r#"time $SLICER write-package "$SOURCE_PATH""#));
216    }
217
218    #[test]
219    fn embedded_script_preserves_catalog_source_revision() {
220        assert!(EMBEDDED_SCRIPT.contains(r#""source_revision": source_revision"#));
221        assert!(EMBEDDED_SCRIPT.contains(r#"variants[variant_name]["source"] = source_entry"#));
222        assert!(EMBEDDED_SCRIPT.contains(r#"existing_variant["source"] = source_entry"#));
223    }
224}
model_package/script.rs

model_package/
script.rs
