1251 lines
45 KiB
Rust
1251 lines
45 KiB
Rust
use serde::de::{MapAccess, Visitor};
|
|
use serde::{Deserialize, Deserializer};
|
|
use serde_json::Value;
|
|
use sha2::{Digest, Sha256};
|
|
use std::collections::{BTreeMap, BTreeSet, HashSet};
|
|
use std::env;
|
|
use std::fmt;
|
|
use std::fs::{self, File, OpenOptions};
|
|
use std::io::{Read, Seek, SeekFrom, Write};
|
|
use std::path::{Path, PathBuf};
|
|
use std::time::Duration;
|
|
|
|
const MAX_SAFETENSORS_HEADER: u64 = 16 * 1024 * 1024;
|
|
const CONVERTER: &str = "qwen38-artifacts-v1-identity";
|
|
const LICENSE: &str = "Qwen Community License 1.0";
|
|
const SOURCE_REPOSITORY: &str = "Youssofal/Qwen3.8-Flash-Next-MTPLX-Bare-Speed";
|
|
const SOURCE_REVISION: &str = "74559cdf34fbfc0b593de72d17e93f37fd4f9ea7";
|
|
const BASE_REPOSITORY: &str = "Qwen/Qwen3.8-Flash-Next";
|
|
const BASE_REVISION: &str = "de4b8e4d43b917e7706784d8bb445c9af86a3540";
|
|
const DOWNLOAD_TIMEOUT: Duration = Duration::from_secs(30);
|
|
|
|
#[derive(Deserialize)]
|
|
struct Manifest {
|
|
format: u32,
|
|
source: Source,
|
|
tensor_inventory: String,
|
|
tensor_inventory_sha256: String,
|
|
files: Vec<Artifact>,
|
|
config: BTreeMap<String, Value>,
|
|
runtime: BTreeMap<String, Value>,
|
|
representative_rows: Vec<Sample>,
|
|
excluded: Vec<Excluded>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct Source {
|
|
repository: String,
|
|
url: String,
|
|
revision: String,
|
|
base_repository: String,
|
|
base_url: String,
|
|
base_revision: String,
|
|
license: String,
|
|
converter: String,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct Artifact {
|
|
path: String,
|
|
role: String,
|
|
size: u64,
|
|
sha256: String,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct Excluded {
|
|
path: String,
|
|
reason: String,
|
|
size: u64,
|
|
sha256: String,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct Sample {
|
|
class: String,
|
|
file: String,
|
|
tensor: String,
|
|
row: u64,
|
|
values: u64,
|
|
bits: Option<u32>,
|
|
group_size: Option<u64>,
|
|
sha256: String,
|
|
}
|
|
|
|
fn main() {
|
|
if let Err(error) = run() {
|
|
eprintln!("qwen38-artifacts: {error}");
|
|
std::process::exit(1);
|
|
}
|
|
}
|
|
|
|
fn run() -> Result<(), String> {
|
|
let args = env::args_os().skip(1).collect::<Vec<_>>();
|
|
match args.as_slice() {
|
|
[command, manifest, source] if command == "verify" => {
|
|
let manifest = Path::new(manifest);
|
|
verify(manifest, Path::new(source))?;
|
|
println!("verified {}", manifest.display());
|
|
Ok(())
|
|
}
|
|
[command, manifest, source, output] if command == "inventory" => {
|
|
let (manifest, _) = load_manifest(Path::new(manifest))?;
|
|
let inventory = build_inventory(&manifest, Path::new(source), false)?;
|
|
fs::write(output, inventory).map_err(|error| error.to_string())
|
|
}
|
|
[command, manifest, source, output] if command == "materialize" => {
|
|
materialize(Path::new(manifest), Path::new(source), Path::new(output))
|
|
}
|
|
[command, manifest, output] if command == "fetch" => {
|
|
fetch(Path::new(manifest), Path::new(output))
|
|
}
|
|
[command, manifest, source] if command == "sample-digests" => {
|
|
let (manifest, _) = load_manifest(Path::new(manifest))?;
|
|
for sample in &manifest.representative_rows {
|
|
println!("{}\t{}", sample.class, sample_digest(Path::new(source), sample)?);
|
|
}
|
|
Ok(())
|
|
}
|
|
_ => Err("usage: qwen38-artifacts fetch MANIFEST OUTPUT | verify MANIFEST SOURCE | inventory MANIFEST SOURCE OUTPUT | sample-digests MANIFEST SOURCE | materialize MANIFEST SOURCE OUTPUT".into()),
|
|
}
|
|
}
|
|
|
|
fn load_manifest(path: &Path) -> Result<(Manifest, PathBuf), String> {
|
|
let bytes = fs::read(path).map_err(|error| format!("{}: {error}", path.display()))?;
|
|
let manifest: Manifest =
|
|
serde_json::from_slice(&bytes).map_err(|error| format!("{}: {error}", path.display()))?;
|
|
if manifest.format != 1 {
|
|
return Err(format!("unsupported manifest format {}", manifest.format));
|
|
}
|
|
validate_source(&manifest.source)?;
|
|
let directory = path
|
|
.parent()
|
|
.ok_or_else(|| format!("{} has no parent directory", path.display()))?;
|
|
Ok((manifest, directory.to_owned()))
|
|
}
|
|
|
|
fn validate_source(source: &Source) -> Result<(), String> {
|
|
for (name, value) in [
|
|
("source repository", source.repository.as_str()),
|
|
("source URL", source.url.as_str()),
|
|
("source revision", source.revision.as_str()),
|
|
("base repository", source.base_repository.as_str()),
|
|
("base URL", source.base_url.as_str()),
|
|
("base revision", source.base_revision.as_str()),
|
|
("license", source.license.as_str()),
|
|
("converter", source.converter.as_str()),
|
|
] {
|
|
if value.is_empty() {
|
|
return Err(format!("{name} is empty"));
|
|
}
|
|
}
|
|
if source.revision.len() != 40 || source.base_revision.len() != 40 {
|
|
return Err("source revisions must be full 40-character commits".into());
|
|
}
|
|
if source.repository != SOURCE_REPOSITORY
|
|
|| source.revision != SOURCE_REVISION
|
|
|| source.base_repository != BASE_REPOSITORY
|
|
|| source.base_revision != BASE_REVISION
|
|
{
|
|
return Err("manifest does not identify the pinned Bare Speed source and base".into());
|
|
}
|
|
if source.url
|
|
!= format!(
|
|
"https://huggingface.co/{}/tree/{}",
|
|
source.repository, source.revision
|
|
)
|
|
|| source.base_url
|
|
!= format!(
|
|
"https://huggingface.co/{}/tree/{}",
|
|
source.base_repository, source.base_revision
|
|
)
|
|
{
|
|
return Err("source URLs do not match their repositories and revisions".into());
|
|
}
|
|
if source.converter != CONVERTER || source.license != LICENSE {
|
|
return Err("manifest converter or license does not match this tool".into());
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn verify(manifest_path: &Path, source: &Path) -> Result<(), String> {
|
|
let (manifest, manifest_directory) = load_manifest(manifest_path)?;
|
|
validate_file_contract(&manifest)?;
|
|
for artifact in &manifest.files {
|
|
verify_artifact(&source.join(&artifact.path), artifact)?;
|
|
}
|
|
validate_json_contract(source, "config.json", &manifest.config)?;
|
|
validate_json_contract(source, "mtplx_runtime.json", &manifest.runtime)?;
|
|
validate_index(&manifest, source)?;
|
|
validate_samples(&manifest, source)?;
|
|
let actual_inventory = build_inventory(&manifest, source, true)?;
|
|
let inventory_path = manifest_directory.join(&manifest.tensor_inventory);
|
|
let expected_inventory = fs::read_to_string(&inventory_path)
|
|
.map_err(|error| format!("{}: {error}", inventory_path.display()))?;
|
|
let inventory_sha256 = sha256_file(&inventory_path)?;
|
|
if inventory_sha256 != manifest.tensor_inventory_sha256 {
|
|
return Err(format!(
|
|
"{} has SHA-256 {inventory_sha256}, expected {}",
|
|
inventory_path.display(),
|
|
manifest.tensor_inventory_sha256
|
|
));
|
|
}
|
|
if actual_inventory != expected_inventory {
|
|
return Err(format!(
|
|
"tensor inventory does not match {}",
|
|
inventory_path.display()
|
|
));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn validate_samples(manifest: &Manifest, source: &Path) -> Result<(), String> {
|
|
let classes = manifest
|
|
.representative_rows
|
|
.iter()
|
|
.map(|sample| sample.class.as_str())
|
|
.collect::<BTreeSet<_>>();
|
|
let required = ["dense", "expert", "gdn", "mtp", "ple", "qsa"]
|
|
.into_iter()
|
|
.collect::<BTreeSet<_>>();
|
|
if classes != required || manifest.representative_rows.len() != required.len() {
|
|
return Err("representative rows must cover dense, expert, GDN, MTP, PLE, and QSA".into());
|
|
}
|
|
let precision = precision_assignments(manifest, source, true)?;
|
|
for sample in &manifest.representative_rows {
|
|
if sample.values == 0 || sample.values > 4096 {
|
|
return Err(format!("{} sample has an invalid width", sample.class));
|
|
}
|
|
let expected_role = match sample.class.as_str() {
|
|
"ple" => "ple",
|
|
"mtp" => "mtp",
|
|
_ => "core",
|
|
};
|
|
if !manifest
|
|
.files
|
|
.iter()
|
|
.any(|artifact| artifact.path == sample.file && artifact.role == expected_role)
|
|
{
|
|
return Err(format!(
|
|
"{} sample is not in its artifact role",
|
|
sample.class
|
|
));
|
|
}
|
|
match precision.get(&sample.tensor) {
|
|
Some(contract)
|
|
if sample.bits == Some(contract.bits)
|
|
&& sample.group_size == Some(contract.group_size) => {}
|
|
None if sample.bits.is_none() && sample.group_size.is_none() => {}
|
|
_ => {
|
|
return Err(format!(
|
|
"{} sample does not match its tensor precision",
|
|
sample.class
|
|
));
|
|
}
|
|
}
|
|
let actual = sample_digest(source, sample)?;
|
|
if actual != sample.sha256 {
|
|
return Err(format!(
|
|
"{} sample has SHA-256 {actual}, expected {}",
|
|
sample.class, sample.sha256
|
|
));
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn sample_digest(source: &Path, sample: &Sample) -> Result<String, String> {
|
|
let path = source.join(&sample.file);
|
|
let (data_start, tensors) = read_safetensors_header(&path, true)?;
|
|
let tensor = tensors
|
|
.get(&sample.tensor)
|
|
.ok_or_else(|| format!("{} is missing {}", path.display(), sample.tensor))?;
|
|
let mut values = Vec::with_capacity(sample.values as usize);
|
|
match tensor.dtype.as_str() {
|
|
"BF16" => {
|
|
if sample.bits.is_some() || sample.group_size.is_some() {
|
|
return Err(format!(
|
|
"{} BF16 sample declares quantization",
|
|
sample.class
|
|
));
|
|
}
|
|
let columns = *tensor
|
|
.shape
|
|
.last()
|
|
.ok_or_else(|| format!("{} has no dimensions", sample.tensor))?;
|
|
if sample.values > columns {
|
|
return Err(format!(
|
|
"{} sample is wider than its tensor row",
|
|
sample.class
|
|
));
|
|
}
|
|
let start = tensor_byte_offset(data_start, tensor, sample.row, columns * 2)?;
|
|
let bytes = read_bytes(&path, start, sample.values * 2)?;
|
|
for bytes in bytes.chunks_exact(2) {
|
|
values.push(f32::from_bits(
|
|
u32::from(u16::from_le_bytes([bytes[0], bytes[1]])) << 16,
|
|
));
|
|
}
|
|
}
|
|
"U32" => {
|
|
let bits = sample
|
|
.bits
|
|
.ok_or_else(|| format!("{} sample has no bit width", sample.class))?;
|
|
let group_size = sample
|
|
.group_size
|
|
.ok_or_else(|| format!("{} sample has no group size", sample.class))?;
|
|
if !matches!(bits, 2 | 4 | 8) || group_size == 0 {
|
|
return Err(format!("{} sample has invalid quantization", sample.class));
|
|
}
|
|
let packed_columns = *tensor
|
|
.shape
|
|
.last()
|
|
.ok_or_else(|| format!("{} has no dimensions", sample.tensor))?;
|
|
let pack = u64::from(32 / bits);
|
|
let columns = packed_columns
|
|
.checked_mul(pack)
|
|
.ok_or_else(|| format!("{} row width overflows", sample.tensor))?;
|
|
if sample.values > columns || !columns.is_multiple_of(group_size) {
|
|
return Err(format!("{} sample has invalid row width", sample.class));
|
|
}
|
|
let base = sample
|
|
.tensor
|
|
.strip_suffix(".weight")
|
|
.ok_or_else(|| format!("{} is not a quantized weight", sample.tensor))?;
|
|
let scales = tensors
|
|
.get(&format!("{base}.scales"))
|
|
.ok_or_else(|| format!("{} has no scales", sample.tensor))?;
|
|
let biases = tensors
|
|
.get(&format!("{base}.biases"))
|
|
.ok_or_else(|| format!("{} has no biases", sample.tensor))?;
|
|
let groups = columns / group_size;
|
|
if scales.dtype != "BF16"
|
|
|| biases.dtype != "BF16"
|
|
|| scales.shape.last() != Some(&groups)
|
|
|| biases.shape != scales.shape
|
|
{
|
|
return Err(format!(
|
|
"{} has incompatible affine parameters",
|
|
sample.tensor
|
|
));
|
|
}
|
|
let packed_start =
|
|
tensor_byte_offset(data_start, tensor, sample.row, packed_columns * 4)?;
|
|
let scale_start = tensor_byte_offset(data_start, scales, sample.row, groups * 2)?;
|
|
let bias_start = tensor_byte_offset(data_start, biases, sample.row, groups * 2)?;
|
|
let packed = read_bytes(&path, packed_start, packed_columns * 4)?;
|
|
let scales = read_bf16(&path, scale_start, groups)?;
|
|
let biases = read_bf16(&path, bias_start, groups)?;
|
|
let mask = (1_u32 << bits) - 1;
|
|
for index in 0..sample.values {
|
|
let word_index = (index / pack) as usize * 4;
|
|
let word = u32::from_le_bytes([
|
|
packed[word_index],
|
|
packed[word_index + 1],
|
|
packed[word_index + 2],
|
|
packed[word_index + 3],
|
|
]);
|
|
let quantized = (word >> ((index % pack) * u64::from(bits))) & mask;
|
|
let group = (index / group_size) as usize;
|
|
values.push(quantized as f32 * scales[group] + biases[group]);
|
|
}
|
|
}
|
|
dtype => return Err(format!("{} sample uses unsupported {dtype}", sample.class)),
|
|
}
|
|
let mut hash = Sha256::new();
|
|
for value in values {
|
|
hash.update(value.to_le_bytes());
|
|
}
|
|
Ok(hex_digest(hash.finalize()))
|
|
}
|
|
|
|
fn tensor_byte_offset(
|
|
data_start: u64,
|
|
tensor: &Tensor,
|
|
row: u64,
|
|
row_bytes: u64,
|
|
) -> Result<u64, String> {
|
|
let offset = row
|
|
.checked_mul(row_bytes)
|
|
.and_then(|offset| tensor.offsets[0].checked_add(offset))
|
|
.and_then(|offset| data_start.checked_add(offset))
|
|
.ok_or_else(|| "sample offset overflows".to_owned())?;
|
|
let end = offset
|
|
.checked_add(row_bytes)
|
|
.ok_or_else(|| "sample range overflows".to_owned())?;
|
|
if end > data_start.saturating_add(tensor.offsets[1]) {
|
|
return Err("sample row exceeds its tensor".into());
|
|
}
|
|
Ok(offset)
|
|
}
|
|
|
|
fn read_bytes(path: &Path, offset: u64, length: u64) -> Result<Vec<u8>, String> {
|
|
let mut file = File::open(path).map_err(|error| format!("{}: {error}", path.display()))?;
|
|
file.seek(SeekFrom::Start(offset))
|
|
.map_err(|error| format!("{}: {error}", path.display()))?;
|
|
let mut bytes = vec![0_u8; length as usize];
|
|
file.read_exact(&mut bytes)
|
|
.map_err(|error| format!("{}: {error}", path.display()))?;
|
|
Ok(bytes)
|
|
}
|
|
|
|
fn read_bf16(path: &Path, offset: u64, count: u64) -> Result<Vec<f32>, String> {
|
|
Ok(read_bytes(path, offset, count * 2)?
|
|
.chunks_exact(2)
|
|
.map(|bytes| f32::from_bits(u32::from(u16::from_le_bytes([bytes[0], bytes[1]])) << 16))
|
|
.collect())
|
|
}
|
|
|
|
fn validate_file_contract(manifest: &Manifest) -> Result<(), String> {
|
|
validate_name(&manifest.tensor_inventory)?;
|
|
if manifest.tensor_inventory_sha256.len() != 64
|
|
|| !manifest
|
|
.tensor_inventory_sha256
|
|
.bytes()
|
|
.all(|byte| byte.is_ascii_hexdigit())
|
|
{
|
|
return Err("tensor inventory has an invalid SHA-256".into());
|
|
}
|
|
let mut paths = HashSet::new();
|
|
let mut roles = BTreeMap::<&str, usize>::new();
|
|
for artifact in &manifest.files {
|
|
validate_name(&artifact.path)?;
|
|
if !paths.insert(artifact.path.as_str()) {
|
|
return Err(format!("duplicate artifact {}", artifact.path));
|
|
}
|
|
if artifact.size == 0
|
|
|| artifact.sha256.len() != 64
|
|
|| !artifact.sha256.bytes().all(|byte| byte.is_ascii_hexdigit())
|
|
{
|
|
return Err(format!("invalid SHA-256 for {}", artifact.path));
|
|
}
|
|
*roles.entry(&artifact.role).or_default() += 1;
|
|
}
|
|
for excluded in &manifest.excluded {
|
|
validate_name(&excluded.path)?;
|
|
if !paths.insert(excluded.path.as_str()) {
|
|
return Err(format!("{} is both included and excluded", excluded.path));
|
|
}
|
|
if excluded.reason.is_empty() {
|
|
return Err(format!("{} has no exclusion reason", excluded.path));
|
|
}
|
|
if excluded.size == 0
|
|
|| excluded.sha256.len() != 64
|
|
|| !excluded.sha256.bytes().all(|byte| byte.is_ascii_hexdigit())
|
|
{
|
|
return Err(format!("{} has invalid exclusion metadata", excluded.path));
|
|
}
|
|
}
|
|
for (role, count) in [("core", 17), ("ple", 1), ("mtp", 1)] {
|
|
if roles.get(role).copied() != Some(count) {
|
|
return Err(format!("manifest must contain {count} {role} artifact(s)"));
|
|
}
|
|
}
|
|
for required in [
|
|
"model.safetensors.index.json",
|
|
"config.json",
|
|
"generation_config.json",
|
|
"mtplx_runtime.json",
|
|
"chat_template.jinja",
|
|
"tokenizer.json",
|
|
"tokenizer_config.json",
|
|
"vocab.json",
|
|
"merges.txt",
|
|
"LICENSE",
|
|
] {
|
|
if !manifest
|
|
.files
|
|
.iter()
|
|
.any(|artifact| artifact.path == required)
|
|
{
|
|
return Err(format!("manifest is missing {required}"));
|
|
}
|
|
}
|
|
if manifest.files.len() != 29 {
|
|
return Err("manifest must contain exactly 29 selected files".into());
|
|
}
|
|
let excluded = manifest
|
|
.excluded
|
|
.iter()
|
|
.map(|entry| entry.path.as_str())
|
|
.collect::<BTreeSet<_>>();
|
|
if excluded
|
|
!= [
|
|
"model-vision.safetensors",
|
|
"preprocessor_config.json",
|
|
"processor_config.json",
|
|
"video_preprocessor_config.json",
|
|
]
|
|
.into_iter()
|
|
.collect()
|
|
{
|
|
return Err("manifest must explicitly exclude the four vision artifacts".into());
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn validate_name(name: &str) -> Result<(), String> {
|
|
let path = Path::new(name);
|
|
if path.components().count() != 1 || path.file_name().is_none() {
|
|
return Err(format!("artifact path must be one file name: {name}"));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn validate_json_contract(
|
|
source: &Path,
|
|
file_name: &str,
|
|
expected_values: &BTreeMap<String, Value>,
|
|
) -> Result<(), String> {
|
|
let path = source.join(file_name);
|
|
let config: Value = serde_json::from_slice(
|
|
&fs::read(&path).map_err(|error| format!("{}: {error}", path.display()))?,
|
|
)
|
|
.map_err(|error| format!("{}: {error}", path.display()))?;
|
|
for (pointer, expected) in expected_values {
|
|
let actual = config
|
|
.pointer(pointer)
|
|
.ok_or_else(|| format!("{file_name} is missing {pointer}"))?;
|
|
if actual != expected {
|
|
return Err(format!(
|
|
"{file_name} {pointer} is {actual}, expected {expected}"
|
|
));
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn validate_index(manifest: &Manifest, source: &Path) -> Result<(), String> {
|
|
let path = source.join("model.safetensors.index.json");
|
|
let value: Value = serde_json::from_slice(
|
|
&fs::read(&path).map_err(|error| format!("{}: {error}", path.display()))?,
|
|
)
|
|
.map_err(|error| format!("{}: {error}", path.display()))?;
|
|
let map = value
|
|
.get("weight_map")
|
|
.and_then(Value::as_object)
|
|
.ok_or_else(|| "model index has no weight_map object".to_owned())?;
|
|
let core_files = manifest
|
|
.files
|
|
.iter()
|
|
.filter(|artifact| artifact.role == "core")
|
|
.map(|artifact| artifact.path.as_str())
|
|
.collect::<HashSet<_>>();
|
|
let indexed = map
|
|
.iter()
|
|
.filter_map(|(name, file)| {
|
|
file.as_str()
|
|
.filter(|file| core_files.contains(file))
|
|
.map(|_| name.clone())
|
|
})
|
|
.collect::<BTreeSet<_>>();
|
|
let actual = tensor_headers(manifest, source, Some("core"), true)?
|
|
.keys()
|
|
.cloned()
|
|
.collect::<BTreeSet<_>>();
|
|
if indexed != actual {
|
|
return Err("model index and core shard tensor inventories differ".into());
|
|
}
|
|
let vision_count = map
|
|
.values()
|
|
.filter(|file| file.as_str() == Some("model-vision.safetensors"))
|
|
.count();
|
|
if vision_count == 0
|
|
|| !manifest
|
|
.excluded
|
|
.iter()
|
|
.any(|entry| entry.path == "model-vision.safetensors")
|
|
{
|
|
return Err("vision tensors are not explicitly excluded".into());
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn build_inventory(manifest: &Manifest, source: &Path, complete: bool) -> Result<String, String> {
|
|
tensor_headers(manifest, source, None, complete)?;
|
|
let precision = precision_assignments(manifest, source, complete)?;
|
|
let mut output = String::from(
|
|
"file\tname\tdtype\tshape\tquant_bits\tgroup_size\tquant_mode\tdata_start\tdata_end\n",
|
|
);
|
|
for artifact in manifest
|
|
.files
|
|
.iter()
|
|
.filter(|artifact| artifact.path.ends_with(".safetensors"))
|
|
{
|
|
let path = source.join(&artifact.path);
|
|
let (data_start, tensors) = read_safetensors_header(&path, complete)?;
|
|
for (name, tensor) in tensors {
|
|
let start = data_start
|
|
.checked_add(tensor.offsets[0])
|
|
.ok_or_else(|| format!("{name} offset overflows"))?;
|
|
let end = data_start
|
|
.checked_add(tensor.offsets[1])
|
|
.ok_or_else(|| format!("{name} offset overflows"))?;
|
|
let quantization = precision.get(&name);
|
|
output.push_str(&format!(
|
|
"{}\t{name}\t{}\t{}\t{}\t{}\t{}\t{start}\t{end}\n",
|
|
artifact.path,
|
|
tensor.dtype,
|
|
tensor
|
|
.shape
|
|
.iter()
|
|
.map(u64::to_string)
|
|
.collect::<Vec<_>>()
|
|
.join("x"),
|
|
quantization.map_or_else(String::new, |value| value.bits.to_string()),
|
|
quantization.map_or_else(String::new, |value| value.group_size.to_string()),
|
|
quantization.map_or("", |value| value.mode.as_str()),
|
|
));
|
|
}
|
|
}
|
|
Ok(output)
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
struct Quantization {
|
|
bits: u32,
|
|
group_size: u64,
|
|
mode: String,
|
|
}
|
|
|
|
fn precision_assignments(
|
|
manifest: &Manifest,
|
|
source: &Path,
|
|
complete: bool,
|
|
) -> Result<BTreeMap<String, Quantization>, String> {
|
|
let core = tensor_headers(manifest, source, Some("core"), complete)?;
|
|
let config_path = source.join("config.json");
|
|
let config: Value = serde_json::from_slice(
|
|
&fs::read(&config_path).map_err(|error| format!("{}: {error}", config_path.display()))?,
|
|
)
|
|
.map_err(|error| format!("{}: {error}", config_path.display()))?;
|
|
if config.get("quantization") != config.get("quantization_config") {
|
|
return Err("config quantization maps differ".into());
|
|
}
|
|
let quantization = config
|
|
.get("quantization")
|
|
.and_then(Value::as_object)
|
|
.ok_or_else(|| "config has no quantization object".to_owned())?;
|
|
let mut assignments = BTreeMap::new();
|
|
for (base, contract) in quantization {
|
|
if matches!(base.as_str(), "bits" | "group_size" | "mode") {
|
|
continue;
|
|
}
|
|
match contract {
|
|
Value::Bool(false) => validate_unquantized(base, &core)?,
|
|
Value::Object(contract) => {
|
|
let bits = contract_u64(contract, "bits", base)? as u32;
|
|
let group_size = contract_u64(contract, "group_size", base)?;
|
|
let mode = contract
|
|
.get("mode")
|
|
.and_then(Value::as_str)
|
|
.ok_or_else(|| format!("{base} has no quantization mode"))?;
|
|
add_quantized(
|
|
base,
|
|
Quantization {
|
|
bits,
|
|
group_size,
|
|
mode: mode.to_owned(),
|
|
},
|
|
&core,
|
|
&mut assignments,
|
|
)?;
|
|
}
|
|
_ => return Err(format!("{base} has an invalid quantization contract")),
|
|
}
|
|
}
|
|
reject_unassigned_quantized(&core, &assignments)?;
|
|
|
|
let mtp = tensor_headers(manifest, source, Some("mtp"), complete)?;
|
|
add_inferred_sidecar(&mtp, 64, &mut assignments)?;
|
|
let ple = tensor_headers(manifest, source, Some("ple"), complete)?;
|
|
let ple_group = config
|
|
.pointer("/mtplx_recipe/ngram/group_size")
|
|
.and_then(Value::as_u64)
|
|
.ok_or_else(|| "config has no PLE group size".to_owned())?;
|
|
let ple_bits = config
|
|
.pointer("/mtplx_recipe/ngram/bits")
|
|
.and_then(Value::as_u64)
|
|
.ok_or_else(|| "config has no PLE bit width".to_owned())? as u32;
|
|
for tensor in ple.keys().filter(|name| name.ends_with(".weight")) {
|
|
let base = tensor.trim_end_matches(".weight");
|
|
add_quantized(
|
|
base,
|
|
Quantization {
|
|
bits: ple_bits,
|
|
group_size: ple_group,
|
|
mode: "affine".into(),
|
|
},
|
|
&ple,
|
|
&mut assignments,
|
|
)?;
|
|
}
|
|
reject_unassigned_quantized(&ple, &assignments)?;
|
|
Ok(assignments)
|
|
}
|
|
|
|
fn contract_u64(
|
|
contract: &serde_json::Map<String, Value>,
|
|
field: &str,
|
|
base: &str,
|
|
) -> Result<u64, String> {
|
|
contract
|
|
.get(field)
|
|
.and_then(Value::as_u64)
|
|
.ok_or_else(|| format!("{base} has no quantization {field}"))
|
|
}
|
|
|
|
fn validate_unquantized(base: &str, tensors: &BTreeMap<String, Tensor>) -> Result<(), String> {
|
|
if let Some(weight) = tensors.get(&format!("{base}.weight"))
|
|
&& weight.dtype != "BF16"
|
|
{
|
|
return Err(format!("{base} is unexpectedly quantized"));
|
|
}
|
|
if tensors.contains_key(&format!("{base}.scales"))
|
|
|| tensors.contains_key(&format!("{base}.biases"))
|
|
{
|
|
return Err(format!("{base} has unexpected affine parameters"));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn add_inferred_sidecar(
|
|
tensors: &BTreeMap<String, Tensor>,
|
|
group_size: u64,
|
|
assignments: &mut BTreeMap<String, Quantization>,
|
|
) -> Result<(), String> {
|
|
for (name, tensor) in tensors {
|
|
if tensor.dtype != "U32" {
|
|
continue;
|
|
}
|
|
let base = name
|
|
.strip_suffix(".weight")
|
|
.ok_or_else(|| format!("quantized sidecar tensor {name} is not a weight"))?;
|
|
let scales = tensors
|
|
.get(&format!("{base}.scales"))
|
|
.ok_or_else(|| format!("{base} has no scales"))?;
|
|
let packed = *tensor
|
|
.shape
|
|
.last()
|
|
.ok_or_else(|| format!("{name} has no dimensions"))?;
|
|
let groups = *scales
|
|
.shape
|
|
.last()
|
|
.ok_or_else(|| format!("{base}.scales has no dimensions"))?;
|
|
let pack = groups
|
|
.checked_mul(group_size)
|
|
.and_then(|columns| columns.checked_div(packed))
|
|
.filter(|pack| matches!(pack, 4 | 8 | 16))
|
|
.ok_or_else(|| format!("{base} has an invalid packed width"))?;
|
|
add_quantized(
|
|
base,
|
|
Quantization {
|
|
bits: (32 / pack) as u32,
|
|
group_size,
|
|
mode: "affine".into(),
|
|
},
|
|
tensors,
|
|
assignments,
|
|
)?;
|
|
}
|
|
reject_unassigned_quantized(tensors, assignments)
|
|
}
|
|
|
|
fn add_quantized(
|
|
base: &str,
|
|
quantization: Quantization,
|
|
tensors: &BTreeMap<String, Tensor>,
|
|
assignments: &mut BTreeMap<String, Quantization>,
|
|
) -> Result<(), String> {
|
|
if !matches!(quantization.bits, 2 | 4 | 8)
|
|
|| quantization.group_size == 0
|
|
|| quantization.mode != "affine"
|
|
{
|
|
return Err(format!("{base} has unsupported quantization"));
|
|
}
|
|
let names = [
|
|
format!("{base}.weight"),
|
|
format!("{base}.scales"),
|
|
format!("{base}.biases"),
|
|
];
|
|
let weight = tensors
|
|
.get(&names[0])
|
|
.ok_or_else(|| format!("{base} is missing its weight"))?;
|
|
let scales = tensors
|
|
.get(&names[1])
|
|
.ok_or_else(|| format!("{base} is missing its scales"))?;
|
|
let biases = tensors
|
|
.get(&names[2])
|
|
.ok_or_else(|| format!("{base} is missing its biases"))?;
|
|
let packed = *weight
|
|
.shape
|
|
.last()
|
|
.ok_or_else(|| format!("{base} has no dimensions"))?;
|
|
let groups = *scales
|
|
.shape
|
|
.last()
|
|
.ok_or_else(|| format!("{base}.scales has no dimensions"))?;
|
|
let pack = u64::from(32 / quantization.bits);
|
|
if weight.dtype != "U32"
|
|
|| scales.dtype != "BF16"
|
|
|| biases.dtype != "BF16"
|
|
|| scales.shape != biases.shape
|
|
|| weight.shape[..weight.shape.len() - 1] != scales.shape[..scales.shape.len() - 1]
|
|
|| packed.checked_mul(pack) != groups.checked_mul(quantization.group_size)
|
|
{
|
|
return Err(format!("{base} does not match its affine contract"));
|
|
}
|
|
for name in names {
|
|
if assignments
|
|
.insert(name.clone(), quantization.clone())
|
|
.is_some()
|
|
{
|
|
return Err(format!("duplicate quantization assignment for {name}"));
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn reject_unassigned_quantized(
|
|
tensors: &BTreeMap<String, Tensor>,
|
|
assignments: &BTreeMap<String, Quantization>,
|
|
) -> Result<(), String> {
|
|
for (name, tensor) in tensors {
|
|
let affine_parameter = name.ends_with(".scales") || name.ends_with(".biases");
|
|
if (tensor.dtype == "U32" || affine_parameter) && !assignments.contains_key(name) {
|
|
return Err(format!("{name} is unexpectedly quantized"));
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn tensor_headers(
|
|
manifest: &Manifest,
|
|
source: &Path,
|
|
role: Option<&str>,
|
|
complete: bool,
|
|
) -> Result<BTreeMap<String, Tensor>, String> {
|
|
let mut all = BTreeMap::new();
|
|
for artifact in manifest.files.iter().filter(|artifact| {
|
|
artifact.path.ends_with(".safetensors") && role.is_none_or(|role| artifact.role == role)
|
|
}) {
|
|
let (_, tensors) = read_safetensors_header(&source.join(&artifact.path), complete)?;
|
|
for (name, tensor) in tensors {
|
|
if all.insert(name.clone(), tensor).is_some() {
|
|
return Err(format!("duplicate tensor {name}"));
|
|
}
|
|
}
|
|
}
|
|
Ok(all)
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct Tensor {
|
|
dtype: String,
|
|
shape: Vec<u64>,
|
|
#[serde(rename = "data_offsets")]
|
|
offsets: [u64; 2],
|
|
}
|
|
|
|
fn read_safetensors_header(
|
|
path: &Path,
|
|
complete: bool,
|
|
) -> Result<(u64, BTreeMap<String, Tensor>), String> {
|
|
let mut file = File::open(path).map_err(|error| format!("{}: {error}", path.display()))?;
|
|
let size = file
|
|
.metadata()
|
|
.map_err(|error| format!("{}: {error}", path.display()))?
|
|
.len();
|
|
let mut length = [0_u8; 8];
|
|
file.read_exact(&mut length)
|
|
.map_err(|error| format!("{}: {error}", path.display()))?;
|
|
let length = u64::from_le_bytes(length);
|
|
if length == 0 || length > MAX_SAFETENSORS_HEADER {
|
|
return Err(format!(
|
|
"{} has invalid header size {length}",
|
|
path.display()
|
|
));
|
|
}
|
|
let data_start = 8_u64
|
|
.checked_add(length)
|
|
.ok_or_else(|| format!("{} header size overflows", path.display()))?;
|
|
let mut bytes = vec![0_u8; length as usize];
|
|
file.read_exact(&mut bytes)
|
|
.map_err(|error| format!("{}: {error}", path.display()))?;
|
|
let mut value = serde_json::from_slice::<UniqueObject>(&bytes)
|
|
.map_err(|error| format!("{}: {error}", path.display()))?;
|
|
value.0.remove("__metadata__");
|
|
let mut tensors = BTreeMap::new();
|
|
for (name, value) in value.0 {
|
|
let tensor: Tensor = serde_json::from_value(value)
|
|
.map_err(|error| format!("{} tensor {name}: {error}", path.display()))?;
|
|
if tensor.shape.is_empty() || tensor.offsets[0] > tensor.offsets[1] {
|
|
return Err(format!(
|
|
"{} tensor {name} has invalid layout",
|
|
path.display()
|
|
));
|
|
}
|
|
let element_size = match tensor.dtype.as_str() {
|
|
"BF16" => 2_u64,
|
|
"I64" => 8_u64,
|
|
"U32" => 4_u64,
|
|
_ => {
|
|
return Err(format!(
|
|
"{} tensor {name} has unsupported dtype",
|
|
path.display()
|
|
));
|
|
}
|
|
};
|
|
let expected_bytes = tensor
|
|
.shape
|
|
.iter()
|
|
.try_fold(element_size, |bytes, dimension| {
|
|
bytes.checked_mul(*dimension)
|
|
});
|
|
if expected_bytes != tensor.offsets[1].checked_sub(tensor.offsets[0]) {
|
|
return Err(format!(
|
|
"{} tensor {name} has a shape/byte-length mismatch",
|
|
path.display()
|
|
));
|
|
}
|
|
if complete && data_start.saturating_add(tensor.offsets[1]) > size {
|
|
return Err(format!("{} tensor {name} exceeds its file", path.display()));
|
|
}
|
|
tensors.insert(name, tensor);
|
|
}
|
|
Ok((data_start, tensors))
|
|
}
|
|
|
|
struct UniqueObject(BTreeMap<String, Value>);
|
|
|
|
impl<'de> Deserialize<'de> for UniqueObject {
|
|
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
|
|
where
|
|
D: Deserializer<'de>,
|
|
{
|
|
struct UniqueObjectVisitor;
|
|
|
|
impl<'de> Visitor<'de> for UniqueObjectVisitor {
|
|
type Value = UniqueObject;
|
|
|
|
fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
formatter.write_str("a safetensors header object with unique tensor names")
|
|
}
|
|
|
|
fn visit_map<M>(self, mut map: M) -> Result<Self::Value, M::Error>
|
|
where
|
|
M: MapAccess<'de>,
|
|
{
|
|
let mut values = BTreeMap::new();
|
|
while let Some((name, value)) = map.next_entry::<String, Value>()? {
|
|
if values.insert(name.clone(), value).is_some() {
|
|
return Err(serde::de::Error::custom(format!(
|
|
"duplicate tensor name {name}"
|
|
)));
|
|
}
|
|
}
|
|
Ok(UniqueObject(values))
|
|
}
|
|
}
|
|
|
|
deserializer.deserialize_map(UniqueObjectVisitor)
|
|
}
|
|
}
|
|
|
|
fn sha256_file(path: &Path) -> Result<String, String> {
|
|
let mut file = File::open(path).map_err(|error| format!("{}: {error}", path.display()))?;
|
|
let mut hash = Sha256::new();
|
|
let mut buffer = vec![0_u8; 8 * 1024 * 1024];
|
|
loop {
|
|
let read = file
|
|
.read(&mut buffer)
|
|
.map_err(|error| format!("{}: {error}", path.display()))?;
|
|
if read == 0 {
|
|
break;
|
|
}
|
|
hash.update(&buffer[..read]);
|
|
}
|
|
Ok(hex_digest(hash.finalize()))
|
|
}
|
|
|
|
fn hex_digest(digest: impl AsRef<[u8]>) -> String {
|
|
let mut output = String::with_capacity(64);
|
|
for byte in digest.as_ref() {
|
|
output.push_str(&format!("{byte:02x}"));
|
|
}
|
|
output
|
|
}
|
|
|
|
fn fetch(manifest_path: &Path, output: &Path) -> Result<(), String> {
|
|
let (manifest, _) = load_manifest(manifest_path)?;
|
|
validate_file_contract(&manifest)?;
|
|
fs::create_dir_all(output).map_err(|error| format!("{}: {error}", output.display()))?;
|
|
for artifact in &manifest.files {
|
|
let destination = output.join(&artifact.path);
|
|
if destination.exists() {
|
|
verify_artifact(&destination, artifact)?;
|
|
continue;
|
|
}
|
|
let partial = output.join(format!("{}.part", artifact.path));
|
|
if partial.metadata().map_or(0, |metadata| metadata.len()) > artifact.size {
|
|
File::create(&partial).map_err(|error| format!("{}: {error}", partial.display()))?;
|
|
}
|
|
if partial.metadata().map_or(0, |metadata| metadata.len()) == artifact.size {
|
|
verify_artifact(&partial, artifact)?;
|
|
fs::rename(&partial, &destination)
|
|
.map_err(|error| format!("{}: {error}", destination.display()))?;
|
|
println!("authenticated {}", artifact.path);
|
|
continue;
|
|
}
|
|
let url = format!(
|
|
"https://huggingface.co/{}/resolve/{}/{}",
|
|
manifest.source.repository, manifest.source.revision, artifact.path
|
|
);
|
|
download(&url, &partial, true)?;
|
|
verify_artifact(&partial, artifact)?;
|
|
fs::rename(&partial, &destination)
|
|
.map_err(|error| format!("{}: {error}", destination.display()))?;
|
|
println!("fetched {}", artifact.path);
|
|
}
|
|
verify(manifest_path, output)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn verify_artifact(path: &Path, artifact: &Artifact) -> Result<(), String> {
|
|
let size = path
|
|
.metadata()
|
|
.map_err(|error| format!("{}: {error}", path.display()))?
|
|
.len();
|
|
if size != artifact.size {
|
|
return Err(format!(
|
|
"{} has {size} bytes, expected {}",
|
|
path.display(),
|
|
artifact.size
|
|
));
|
|
}
|
|
let actual = sha256_file(path)?;
|
|
if actual != artifact.sha256 {
|
|
return Err(format!(
|
|
"{} has SHA-256 {actual}, expected {}",
|
|
path.display(),
|
|
artifact.sha256
|
|
));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn download(url: &str, partial: &Path, https_only: bool) -> Result<(), String> {
|
|
let offset = partial.metadata().map_or(0, |metadata| metadata.len());
|
|
let agent: ureq::Agent = ureq::Agent::config_builder()
|
|
.https_only(https_only)
|
|
.timeout_connect(Some(DOWNLOAD_TIMEOUT))
|
|
.timeout_recv_response(Some(DOWNLOAD_TIMEOUT))
|
|
.timeout_recv_body(Some(DOWNLOAD_TIMEOUT))
|
|
.build()
|
|
.into();
|
|
let mut request = agent.get(url);
|
|
if offset > 0 {
|
|
request = request.header("Range", format!("bytes={offset}-"));
|
|
}
|
|
let mut response = request
|
|
.call()
|
|
.map_err(|error| format!("artifact download failed: {error}"))?;
|
|
let status = response.status().as_u16();
|
|
let append = offset > 0 && status == 206;
|
|
if offset > 0 && !matches!(status, 200 | 206) {
|
|
return Err(format!(
|
|
"artifact server returned HTTP {status} while resuming at byte {offset}"
|
|
));
|
|
}
|
|
if append {
|
|
let expected = format!("bytes {offset}-");
|
|
let content_range = response
|
|
.headers()
|
|
.get("content-range")
|
|
.and_then(|value| value.to_str().ok())
|
|
.unwrap_or_default();
|
|
if !content_range.starts_with(&expected) {
|
|
return Err(format!(
|
|
"artifact server returned an invalid Content-Range while resuming at byte {offset}"
|
|
));
|
|
}
|
|
}
|
|
let mut file = OpenOptions::new()
|
|
.create(true)
|
|
.write(true)
|
|
.append(append)
|
|
.truncate(!append)
|
|
.open(partial)
|
|
.map_err(|error| format!("{}: {error}", partial.display()))?;
|
|
let mut body = response.body_mut().as_reader();
|
|
let mut buffer = vec![0_u8; 1024 * 1024];
|
|
loop {
|
|
let read = body
|
|
.read(&mut buffer)
|
|
.map_err(|error| format!("artifact download failed: {error}"))?;
|
|
if read == 0 {
|
|
break;
|
|
}
|
|
file.write_all(&buffer[..read])
|
|
.map_err(|error| format!("{}: {error}", partial.display()))?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn materialize(manifest_path: &Path, source: &Path, output: &Path) -> Result<(), String> {
|
|
verify(manifest_path, source)?;
|
|
fs::create_dir(output).map_err(|error| format!("{}: {error}", output.display()))?;
|
|
let (manifest, manifest_directory) = load_manifest(manifest_path)?;
|
|
for artifact in &manifest.files {
|
|
let from = source.join(&artifact.path);
|
|
let to = output.join(&artifact.path);
|
|
if fs::hard_link(&from, &to).is_err() {
|
|
fs::copy(&from, &to).map_err(|error| format!("{}: {error}", to.display()))?;
|
|
}
|
|
}
|
|
fs::copy(manifest_path, output.join("manifest.json"))
|
|
.map_err(|error| format!("{}: {error}", output.display()))?;
|
|
fs::copy(
|
|
manifest_directory.join(&manifest.tensor_inventory),
|
|
output.join(&manifest.tensor_inventory),
|
|
)
|
|
.map_err(|error| format!("{}: {error}", output.display()))?;
|
|
verify(manifest_path, output)?;
|
|
println!(
|
|
"materialized {} files ({} bytes) in {}",
|
|
manifest.files.len(),
|
|
manifest
|
|
.files
|
|
.iter()
|
|
.map(|artifact| artifact.size)
|
|
.sum::<u64>(),
|
|
output.display()
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::net::TcpListener;
|
|
use std::thread;
|
|
|
|
fn safetensors_fixture(name: &str, header: &str, payload: &[u8]) -> PathBuf {
|
|
let path = env::temp_dir().join(format!(
|
|
"ds4-qwen38-{name}-{}-{}.safetensors",
|
|
std::process::id(),
|
|
std::thread::current().name().unwrap_or("test")
|
|
));
|
|
let mut bytes = (header.len() as u64).to_le_bytes().to_vec();
|
|
bytes.extend_from_slice(header.as_bytes());
|
|
bytes.extend_from_slice(payload);
|
|
fs::write(&path, bytes).unwrap();
|
|
path
|
|
}
|
|
|
|
#[test]
|
|
fn artifact_names_cannot_escape_the_output_directory() {
|
|
assert!(validate_name("model.safetensors").is_ok());
|
|
assert!(validate_name("../model.safetensors").is_err());
|
|
assert!(validate_name("nested/model.safetensors").is_err());
|
|
assert!(validate_name("").is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn safetensors_validation_rejects_duplicate_names_and_invalid_ranges() {
|
|
let duplicate = safetensors_fixture(
|
|
"duplicate",
|
|
r#"{"x":{"dtype":"BF16","shape":[1],"data_offsets":[0,2]},"x":{"dtype":"BF16","shape":[1],"data_offsets":[0,2]}}"#,
|
|
&[0, 0],
|
|
);
|
|
assert!(read_safetensors_header(&duplicate, true).is_err());
|
|
fs::remove_file(duplicate).unwrap();
|
|
|
|
let outside = safetensors_fixture(
|
|
"outside",
|
|
r#"{"x":{"dtype":"BF16","shape":[1],"data_offsets":[0,2]}}"#,
|
|
&[0],
|
|
);
|
|
assert!(read_safetensors_header(&outside, true).is_err());
|
|
fs::remove_file(outside).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn affine_sample_digest_uses_packed_values_scales_and_biases() {
|
|
let header = r#"{"x.biases":{"dtype":"BF16","shape":[1,1],"data_offsets":[6,8]},"x.scales":{"dtype":"BF16","shape":[1,1],"data_offsets":[4,6]},"x.weight":{"dtype":"U32","shape":[1,1],"data_offsets":[0,4]}}"#;
|
|
let mut payload = 0x7654_3210_u32.to_le_bytes().to_vec();
|
|
payload.extend_from_slice(&0x4000_u16.to_le_bytes());
|
|
payload.extend_from_slice(&0x3f80_u16.to_le_bytes());
|
|
let path = safetensors_fixture("affine", header, &payload);
|
|
let sample = Sample {
|
|
class: "fixture".into(),
|
|
file: path.file_name().unwrap().to_string_lossy().into_owned(),
|
|
tensor: "x.weight".into(),
|
|
row: 0,
|
|
values: 8,
|
|
bits: Some(4),
|
|
group_size: Some(8),
|
|
sha256: String::new(),
|
|
};
|
|
let mut expected = Sha256::new();
|
|
for value in [1.0_f32, 3.0, 5.0, 7.0, 9.0, 11.0, 13.0, 15.0] {
|
|
expected.update(value.to_le_bytes());
|
|
}
|
|
assert_eq!(
|
|
sample_digest(path.parent().unwrap(), &sample).unwrap(),
|
|
hex_digest(expected.finalize())
|
|
);
|
|
fs::remove_file(path).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn artifact_download_resumes_an_authenticated_partial() {
|
|
let content = b"authenticated artifact";
|
|
let offset = 7;
|
|
let directory = env::temp_dir().join(format!("ds4-qwen38-fetch-{}", std::process::id()));
|
|
fs::create_dir_all(&directory).unwrap();
|
|
let partial = directory.join("artifact.part");
|
|
fs::write(&partial, &content[..offset]).unwrap();
|
|
let listener = match TcpListener::bind("127.0.0.1:0") {
|
|
Ok(listener) => listener,
|
|
Err(error) if error.kind() == std::io::ErrorKind::PermissionDenied => {
|
|
fs::remove_dir_all(directory).unwrap();
|
|
return;
|
|
}
|
|
Err(error) => panic!("could not start test server: {error}"),
|
|
};
|
|
let address = listener.local_addr().unwrap();
|
|
let server = thread::spawn(move || {
|
|
let (mut connection, _) = listener.accept().unwrap();
|
|
let mut request = [0_u8; 2048];
|
|
let read = connection.read(&mut request).unwrap();
|
|
let request = String::from_utf8_lossy(&request[..read]).to_ascii_lowercase();
|
|
assert!(request.contains("range: bytes=7-"));
|
|
let remaining = &content[offset..];
|
|
write!(
|
|
connection,
|
|
"HTTP/1.1 206 Partial Content\r\nContent-Length: {}\r\nContent-Range: bytes {offset}-{}/{}\r\nConnection: close\r\n\r\n",
|
|
remaining.len(),
|
|
content.len() - 1,
|
|
content.len(),
|
|
)
|
|
.unwrap();
|
|
connection.write_all(remaining).unwrap();
|
|
});
|
|
|
|
download(&format!("http://{address}/artifact"), &partial, false).unwrap();
|
|
server.join().unwrap();
|
|
verify_artifact(
|
|
&partial,
|
|
&Artifact {
|
|
path: "artifact".into(),
|
|
role: "fixture".into(),
|
|
size: content.len() as u64,
|
|
sha256: "ef056ce05d5f86154a8498b5504a46bc4b633fb68ec795a3a1cfd0548f791e4c".into(),
|
|
},
|
|
)
|
|
.unwrap();
|
|
fs::remove_dir_all(directory).unwrap();
|
|
}
|
|
}
|