943 lines
33 KiB
Rust
943 lines
33 KiB
Rust
use super::*;
|
|
use sha2::{Digest, Sha256};
|
|
use std::fs::{self, File, OpenOptions};
|
|
use std::io::{Read, Write};
|
|
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
|
use std::time::Duration;
|
|
|
|
const DOWNLOAD_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
|
|
const DOWNLOAD_STALL_TIMEOUT: Duration = Duration::from_secs(30);
|
|
|
|
pub(crate) fn download_managed_artifact(
|
|
id: ManagedArtifactId,
|
|
models_path: &Path,
|
|
cancel: &AtomicBool,
|
|
verified_bytes: &AtomicU64,
|
|
) -> Result<DownloadOutcome, String> {
|
|
if id.is_qwen() {
|
|
return download_qwen(id, models_path, cancel, verified_bytes);
|
|
}
|
|
download_artifact_with_cancel(
|
|
id.model(),
|
|
id.artifact(),
|
|
models_path,
|
|
cancel,
|
|
verified_bytes,
|
|
)
|
|
}
|
|
|
|
pub(crate) fn validate_managed_artifact(
|
|
id: ManagedArtifactId,
|
|
models_path: &Path,
|
|
cancel: &AtomicBool,
|
|
verified_bytes: &AtomicU64,
|
|
) -> Result<DownloadOutcome, String> {
|
|
if id.is_qwen() {
|
|
return validate_qwen(id, models_path, cancel, verified_bytes);
|
|
}
|
|
let model = id.model();
|
|
let artifact = id.artifact();
|
|
let destination = artifact.path(model, models_path);
|
|
let partial = artifact.partial_path(model, models_path);
|
|
let (path, promote) = if destination.exists() {
|
|
(destination.clone(), false)
|
|
} else if partial.exists() {
|
|
(partial.clone(), true)
|
|
} else {
|
|
return Err(format!("{} is not downloaded", artifact.label));
|
|
};
|
|
|
|
match verify(&path, artifact, model, cancel, verified_bytes) {
|
|
Ok(DownloadOutcome::Stopped) => return Ok(DownloadOutcome::Stopped),
|
|
Ok(DownloadOutcome::Complete) => {}
|
|
Err(error) => {
|
|
let marker = artifact.verification_path(model, models_path);
|
|
if let Err(remove_error) = fs::remove_file(marker)
|
|
&& remove_error.kind() != std::io::ErrorKind::NotFound
|
|
{
|
|
return Err(format!(
|
|
"{error}; could not remove checksum marker: {remove_error}"
|
|
));
|
|
}
|
|
return Err(error);
|
|
}
|
|
}
|
|
if promote {
|
|
fs::rename(partial, destination).map_err(|error| error.to_string())?;
|
|
}
|
|
mark_verified(model, artifact, models_path)?;
|
|
Ok(DownloadOutcome::Complete)
|
|
}
|
|
|
|
pub(crate) fn delete_managed_artifact(
|
|
id: ManagedArtifactId,
|
|
models_path: &Path,
|
|
) -> Result<(), String> {
|
|
if id.is_qwen() {
|
|
return delete_qwen(id, models_path);
|
|
}
|
|
let model = id.model();
|
|
let artifact = id.artifact();
|
|
for path in [
|
|
artifact.path(model, models_path),
|
|
artifact.partial_path(model, models_path),
|
|
artifact.verification_path(model, models_path),
|
|
] {
|
|
match fs::remove_file(path) {
|
|
Ok(()) => {}
|
|
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
|
|
Err(error) => return Err(error.to_string()),
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn download_qwen(
|
|
id: ManagedArtifactId,
|
|
models_path: &Path,
|
|
cancel: &AtomicBool,
|
|
verified_bytes: &AtomicU64,
|
|
) -> Result<DownloadOutcome, String> {
|
|
let root = qwen::root(models_path);
|
|
fs::create_dir_all(&root).map_err(|error| format!("{}: {error}", root.display()))?;
|
|
verified_bytes.store(0, Ordering::Relaxed);
|
|
for artifact in qwen::selected(id) {
|
|
if cancel.load(Ordering::Relaxed) {
|
|
return Ok(DownloadOutcome::Stopped);
|
|
}
|
|
if qwen::artifact_is_installed(&root, artifact) {
|
|
verified_bytes.fetch_add(artifact.size, Ordering::Relaxed);
|
|
continue;
|
|
}
|
|
let destination = qwen::path(&root, artifact);
|
|
let partial = qwen::partial_path(&root, artifact);
|
|
let source = if destination.exists() {
|
|
destination.clone()
|
|
} else {
|
|
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
|
|
&& download_url_to_partial(&qwen::url(artifact), &partial, cancel)?
|
|
== DownloadOutcome::Stopped
|
|
{
|
|
return Ok(DownloadOutcome::Stopped);
|
|
}
|
|
partial.clone()
|
|
};
|
|
match qwen::verify(&source, artifact, cancel, verified_bytes) {
|
|
Ok(false) => return Ok(DownloadOutcome::Stopped),
|
|
Ok(true) => {}
|
|
Err(error) => {
|
|
let _ = fs::remove_file(qwen::verification_path(&root, artifact));
|
|
if source == partial {
|
|
fs::remove_file(&partial).map_err(|remove_error| {
|
|
format!("{error}; could not remove partial file: {remove_error}")
|
|
})?;
|
|
}
|
|
return Err(error);
|
|
}
|
|
}
|
|
if source == partial {
|
|
fs::rename(&partial, &destination)
|
|
.map_err(|error| format!("{}: {error}", destination.display()))?;
|
|
}
|
|
fs::write(qwen::verification_path(&root, artifact), &artifact.sha256)
|
|
.map_err(|error| error.to_string())?;
|
|
}
|
|
if id == ManagedArtifactId::Qwen38FlashNext
|
|
&& let Err(error) = crate::engine::validate_qwen_artifacts(&root)
|
|
{
|
|
clear_qwen_markers(id, &root)?;
|
|
return Err(error);
|
|
}
|
|
Ok(DownloadOutcome::Complete)
|
|
}
|
|
|
|
fn validate_qwen(
|
|
id: ManagedArtifactId,
|
|
models_path: &Path,
|
|
cancel: &AtomicBool,
|
|
verified_bytes: &AtomicU64,
|
|
) -> Result<DownloadOutcome, String> {
|
|
let root = qwen::root(models_path);
|
|
verified_bytes.store(0, Ordering::Relaxed);
|
|
for artifact in qwen::selected(id) {
|
|
if !qwen::path(&root, artifact).exists() && !qwen::partial_path(&root, artifact).exists() {
|
|
return Err(format!("{} is not downloaded", artifact.path));
|
|
}
|
|
}
|
|
for artifact in qwen::selected(id) {
|
|
let destination = qwen::path(&root, artifact);
|
|
let partial = qwen::partial_path(&root, artifact);
|
|
let (source, promote) = if destination.exists() {
|
|
(destination.clone(), false)
|
|
} else if partial.exists() {
|
|
(partial.clone(), true)
|
|
} else {
|
|
return Err(format!("{} is not downloaded", artifact.path));
|
|
};
|
|
match qwen::verify(&source, artifact, cancel, verified_bytes) {
|
|
Ok(false) => return Ok(DownloadOutcome::Stopped),
|
|
Ok(true) => {}
|
|
Err(error) => {
|
|
let _ = fs::remove_file(qwen::verification_path(&root, artifact));
|
|
return Err(error);
|
|
}
|
|
}
|
|
if promote {
|
|
fs::rename(&partial, &destination)
|
|
.map_err(|error| format!("{}: {error}", destination.display()))?;
|
|
}
|
|
fs::write(qwen::verification_path(&root, artifact), &artifact.sha256)
|
|
.map_err(|error| error.to_string())?;
|
|
}
|
|
if id == ManagedArtifactId::Qwen38FlashNext
|
|
&& let Err(error) = crate::engine::validate_qwen_artifacts(&root)
|
|
{
|
|
clear_qwen_markers(id, &root)?;
|
|
return Err(error);
|
|
}
|
|
Ok(DownloadOutcome::Complete)
|
|
}
|
|
|
|
fn clear_qwen_markers(id: ManagedArtifactId, root: &Path) -> Result<(), String> {
|
|
for artifact in qwen::selected(id) {
|
|
let path = qwen::verification_path(root, artifact);
|
|
match fs::remove_file(&path) {
|
|
Ok(()) => {}
|
|
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
|
|
Err(error) => return Err(format!("{}: {error}", path.display())),
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn delete_qwen(id: ManagedArtifactId, models_path: &Path) -> Result<(), String> {
|
|
let root = qwen::root(models_path);
|
|
for artifact in qwen::selected(id) {
|
|
for path in [
|
|
qwen::path(&root, artifact),
|
|
qwen::partial_path(&root, artifact),
|
|
qwen::verification_path(&root, artifact),
|
|
] {
|
|
match fs::remove_file(&path) {
|
|
Ok(()) => {}
|
|
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
|
|
Err(error) => return Err(format!("{}: {error}", path.display())),
|
|
}
|
|
}
|
|
}
|
|
match fs::remove_dir(&root) {
|
|
Ok(()) => Ok(()),
|
|
Err(error)
|
|
if matches!(
|
|
error.kind(),
|
|
std::io::ErrorKind::NotFound | std::io::ErrorKind::DirectoryNotEmpty
|
|
) =>
|
|
{
|
|
Ok(())
|
|
}
|
|
Err(error) => Err(format!("{}: {error}", root.display())),
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn download_artifact(
|
|
model: ModelChoice,
|
|
artifact: &Artifact,
|
|
models_path: &Path,
|
|
) -> Result<DownloadOutcome, String> {
|
|
download_artifact_with_cancel(
|
|
model,
|
|
artifact,
|
|
models_path,
|
|
&AtomicBool::new(false),
|
|
&AtomicU64::new(0),
|
|
)
|
|
}
|
|
|
|
fn download_artifact_with_cancel(
|
|
model: ModelChoice,
|
|
artifact: &Artifact,
|
|
models_path: &Path,
|
|
cancel: &AtomicBool,
|
|
verified_bytes: &AtomicU64,
|
|
) -> Result<DownloadOutcome, String> {
|
|
if cancel.load(Ordering::Relaxed) {
|
|
return Ok(DownloadOutcome::Stopped);
|
|
}
|
|
let destination = artifact.path(model, models_path);
|
|
if artifact.is_installed(model, models_path) {
|
|
return Ok(DownloadOutcome::Complete);
|
|
}
|
|
if destination.exists() {
|
|
if verify(&destination, artifact, model, cancel, verified_bytes)?
|
|
== DownloadOutcome::Stopped
|
|
{
|
|
return Ok(DownloadOutcome::Stopped);
|
|
}
|
|
mark_verified(model, artifact, models_path)?;
|
|
return Ok(DownloadOutcome::Complete);
|
|
}
|
|
|
|
let directory = destination
|
|
.parent()
|
|
.ok_or_else(|| "model artifact path has no parent directory".to_owned())?;
|
|
fs::create_dir_all(directory).map_err(|error| error.to_string())?;
|
|
let partial = artifact.partial_path(model, models_path);
|
|
let partial_size = partial.metadata().map_or(0, |metadata| metadata.len());
|
|
if partial_size > artifact.size {
|
|
File::create(&partial).map_err(|error| error.to_string())?;
|
|
}
|
|
if partial.metadata().map_or(0, |metadata| metadata.len()) != artifact.size
|
|
&& download_to_partial(artifact, &partial, cancel)? == DownloadOutcome::Stopped
|
|
{
|
|
return Ok(DownloadOutcome::Stopped);
|
|
}
|
|
|
|
if cancel.load(Ordering::Relaxed) {
|
|
return Ok(DownloadOutcome::Stopped);
|
|
}
|
|
match verify(&partial, artifact, model, cancel, verified_bytes) {
|
|
Ok(DownloadOutcome::Stopped) => return Ok(DownloadOutcome::Stopped),
|
|
Ok(DownloadOutcome::Complete) => {}
|
|
Err(error) => {
|
|
fs::remove_file(&partial).map_err(|remove_error| {
|
|
format!("{error}; could not remove partial file: {remove_error}")
|
|
})?;
|
|
return Err(error);
|
|
}
|
|
}
|
|
fs::rename(partial, destination).map_err(|error| error.to_string())?;
|
|
mark_verified(model, artifact, models_path)?;
|
|
Ok(DownloadOutcome::Complete)
|
|
}
|
|
|
|
fn mark_verified(
|
|
model: ModelChoice,
|
|
artifact: &Artifact,
|
|
models_path: &Path,
|
|
) -> Result<(), String> {
|
|
fs::write(
|
|
artifact.verification_path(model, models_path),
|
|
artifact.sha256,
|
|
)
|
|
.map_err(|error| error.to_string())
|
|
}
|
|
|
|
fn verify(
|
|
path: &Path,
|
|
artifact: &Artifact,
|
|
model: ModelChoice,
|
|
cancel: &AtomicBool,
|
|
verified_bytes: &AtomicU64,
|
|
) -> Result<DownloadOutcome, String> {
|
|
verified_bytes.store(0, Ordering::Relaxed);
|
|
let size = path.metadata().map_err(|error| error.to_string())?.len();
|
|
if size != artifact.size {
|
|
return Err(format!(
|
|
"{} has size {size}, expected {}",
|
|
path.display(),
|
|
artifact.size
|
|
));
|
|
}
|
|
|
|
let mut file = File::open(path).map_err(|error| error.to_string())?;
|
|
let mut hasher = Sha256::new();
|
|
let mut buffer = vec![0; 1024 * 1024];
|
|
loop {
|
|
if cancel.load(Ordering::Relaxed) {
|
|
return Ok(DownloadOutcome::Stopped);
|
|
}
|
|
let count = file.read(&mut buffer).map_err(|error| error.to_string())?;
|
|
if count == 0 {
|
|
break;
|
|
}
|
|
hasher.update(&buffer[..count]);
|
|
verified_bytes.fetch_add(count as u64, Ordering::Relaxed);
|
|
}
|
|
let actual = hex(&hasher.finalize());
|
|
if actual != artifact.sha256 {
|
|
return Err(format!(
|
|
"Checksum verification failed for {}",
|
|
path.display()
|
|
));
|
|
}
|
|
if artifact.vision {
|
|
crate::engine::validate_vision_artifact(path)?;
|
|
} else if let Some(support) = artifact.support {
|
|
crate::engine::validate_model_artifact(path, model, support)?;
|
|
}
|
|
Ok(DownloadOutcome::Complete)
|
|
}
|
|
|
|
fn hex(bytes: &[u8]) -> String {
|
|
const DIGITS: &[u8; 16] = b"0123456789abcdef";
|
|
let mut result = String::with_capacity(bytes.len() * 2);
|
|
for byte in bytes {
|
|
result.push(DIGITS[(byte >> 4) as usize] as char);
|
|
result.push(DIGITS[(byte & 0x0f) as usize] as char);
|
|
}
|
|
result
|
|
}
|
|
|
|
fn download_to_partial(
|
|
artifact: &Artifact,
|
|
partial: &Path,
|
|
cancel: &AtomicBool,
|
|
) -> Result<DownloadOutcome, String> {
|
|
download_url_to_partial(&artifact.url(), partial, cancel)
|
|
}
|
|
|
|
fn download_url_to_partial(
|
|
url: &str,
|
|
partial: &Path,
|
|
cancel: &AtomicBool,
|
|
) -> Result<DownloadOutcome, String> {
|
|
download_url_to_partial_with_timeout(url, partial, cancel, true, DOWNLOAD_STALL_TIMEOUT)
|
|
}
|
|
|
|
fn download_url_to_partial_with_timeout(
|
|
url: &str,
|
|
partial: &Path,
|
|
cancel: &AtomicBool,
|
|
https_only: bool,
|
|
stall_timeout: Duration,
|
|
) -> Result<DownloadOutcome, 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_CONNECT_TIMEOUT))
|
|
.timeout_recv_response(Some(stall_timeout))
|
|
.timeout_recv_body(Some(stall_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!("Model download failed: {error}"))?;
|
|
let status = response.status().as_u16();
|
|
let append = offset > 0 && status == 206;
|
|
if offset > 0 && status != 200 && status != 206 {
|
|
return Err(format!(
|
|
"Model server returned HTTP {status} while resuming at byte {offset}"
|
|
));
|
|
}
|
|
if append {
|
|
validate_content_range(&response, offset)?;
|
|
}
|
|
|
|
let mut output = OpenOptions::new()
|
|
.create(true)
|
|
.write(true)
|
|
.append(append)
|
|
.truncate(!append)
|
|
.open(partial)
|
|
.map_err(|error| error.to_string())?;
|
|
let mut body = response.body_mut().as_reader();
|
|
let mut buffer = vec![0; 1024 * 1024];
|
|
loop {
|
|
if cancel.load(Ordering::Relaxed) {
|
|
return Ok(DownloadOutcome::Stopped);
|
|
}
|
|
let count = body
|
|
.read(&mut buffer)
|
|
.map_err(|error| format!("Model download failed: {error}"))?;
|
|
if count == 0 {
|
|
break;
|
|
}
|
|
output
|
|
.write_all(&buffer[..count])
|
|
.map_err(|error| error.to_string())?;
|
|
}
|
|
Ok(DownloadOutcome::Complete)
|
|
}
|
|
|
|
fn validate_content_range(
|
|
response: &ureq::http::Response<ureq::Body>,
|
|
offset: u64,
|
|
) -> Result<(), String> {
|
|
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) {
|
|
Ok(())
|
|
} else {
|
|
Err(format!(
|
|
"Model server returned an invalid Content-Range while resuming at byte {offset}"
|
|
))
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::net::TcpListener;
|
|
use std::thread;
|
|
use std::time::{SystemTime, UNIX_EPOCH};
|
|
|
|
#[test]
|
|
fn catalog_and_checksum_verification_are_explicit() {
|
|
assert_eq!(ModelChoice::from_id("glm-5.2"), Some(ModelChoice::Glm52));
|
|
assert_eq!(
|
|
ModelChoice::from_id("deepseek-v4-flash-0731"),
|
|
Some(ModelChoice::DeepSeekV4Flash0731)
|
|
);
|
|
assert!(ModelChoice::from_id("unknown").is_none());
|
|
assert_eq!(MODEL_CHOICES.len(), 6);
|
|
assert_eq!(MANAGED_ARTIFACTS.len(), 11);
|
|
assert_eq!(
|
|
ModelChoice::from_id("qwen3.8-flash-next"),
|
|
Some(ModelChoice::Qwen38FlashNext)
|
|
);
|
|
assert_eq!(ModelChoice::Glm52.main_artifact().size, 211_075_856_448);
|
|
assert_eq!(
|
|
ModelChoice::DeepSeekV4Flash0731.main_artifact().size,
|
|
86_720_111_488
|
|
);
|
|
assert_eq!(ModelChoice::DeepSeekV4Flash0731.artifacts(true).count(), 2);
|
|
assert_eq!(ModelChoice::Glm52.artifacts(true).count(), 1);
|
|
assert_eq!(ModelChoice::Glm53Flash.artifacts(true).count(), 2);
|
|
assert_eq!(
|
|
engine_artifacts(ModelChoice::Qwen38FlashNext, false, Path::new("/models")).model,
|
|
Path::new("/models/qwen3.8-flash-next")
|
|
);
|
|
|
|
let id = SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_nanos();
|
|
let models_path = std::env::temp_dir().join(format!("ds4-server-models-{id}"));
|
|
let engine = engine_artifacts(ModelChoice::DeepSeekV4Flash0731, true, &models_path);
|
|
assert_eq!(
|
|
engine.model.file_name(),
|
|
Some(std::ffi::OsStr::new(FLASH_0731.file_name))
|
|
);
|
|
assert_eq!(
|
|
engine.support.as_deref().and_then(Path::file_name),
|
|
Some(std::ffi::OsStr::new(FLASH_0731_DSPARK.file_name))
|
|
);
|
|
assert!(
|
|
engine_artifacts(ModelChoice::DeepSeekV4Flash0731, false, &models_path)
|
|
.support
|
|
.is_none()
|
|
);
|
|
let empty = Artifact {
|
|
label: "empty",
|
|
file_name: "empty",
|
|
repository: "",
|
|
size: 0,
|
|
sha256: "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
|
|
support: None,
|
|
vision: false,
|
|
};
|
|
let partial = empty.partial_path(ModelChoice::DeepSeekV4Flash0731, &models_path);
|
|
fs::create_dir_all(partial.parent().unwrap()).unwrap();
|
|
fs::write(&partial, []).unwrap();
|
|
download_artifact(ModelChoice::DeepSeekV4Flash0731, &empty, &models_path).unwrap();
|
|
assert!(empty.is_installed(ModelChoice::DeepSeekV4Flash0731, &models_path));
|
|
assert!(!partial.exists());
|
|
assert_eq!(
|
|
fs::read_to_string(
|
|
empty.verification_path(ModelChoice::DeepSeekV4Flash0731, &models_path)
|
|
)
|
|
.unwrap(),
|
|
empty.sha256
|
|
);
|
|
let installed = empty.path(ModelChoice::DeepSeekV4Flash0731, &models_path);
|
|
assert!(empty.validate_installed_path(&installed).is_ok());
|
|
let wrong_name = installed.with_file_name("wrong-checkpoint.gguf");
|
|
fs::write(&wrong_name, []).unwrap();
|
|
fs::write(wrong_name.with_extension("gguf.sha256"), empty.sha256).unwrap();
|
|
assert!(empty.validate_installed_path(&wrong_name).is_err());
|
|
assert!(
|
|
validate_engine_artifacts(
|
|
ModelChoice::DeepSeekV4Flash0731,
|
|
true,
|
|
&EngineArtifacts {
|
|
model: installed.clone(),
|
|
support: None,
|
|
vision: None,
|
|
},
|
|
)
|
|
.is_err()
|
|
);
|
|
assert!(
|
|
validate_engine_artifacts(
|
|
ModelChoice::DeepSeekV4Pro,
|
|
true,
|
|
&EngineArtifacts {
|
|
model: installed,
|
|
support: None,
|
|
vision: None,
|
|
},
|
|
)
|
|
.is_err()
|
|
);
|
|
fs::remove_dir_all(models_path).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "requires DS4_QWEN38_ARTIFACTS to point at the pinned 105 GB source"]
|
|
fn qwen_set_verifies_reloads_and_removes_as_one_managed_artifact() {
|
|
let source = std::env::var_os("DS4_QWEN38_ARTIFACTS")
|
|
.map(std::path::PathBuf::from)
|
|
.expect("DS4_QWEN38_ARTIFACTS is set");
|
|
let models = std::env::temp_dir().join(format!(
|
|
"ds4-qwen-managed-{}-{}",
|
|
std::process::id(),
|
|
SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_nanos()
|
|
));
|
|
let root = qwen::root(&models);
|
|
fs::create_dir_all(&root).unwrap();
|
|
for artifact in qwen::artifacts() {
|
|
fs::hard_link(source.join(&artifact.path), qwen::path(&root, artifact)).unwrap();
|
|
}
|
|
fs::remove_file(qwen::path(
|
|
&root,
|
|
qwen::artifacts()
|
|
.iter()
|
|
.find(|artifact| artifact.role == "mtp")
|
|
.unwrap(),
|
|
))
|
|
.unwrap();
|
|
assert!(
|
|
validate_qwen(
|
|
ManagedArtifactId::Qwen38FlashNext,
|
|
&models,
|
|
&AtomicBool::new(false),
|
|
&AtomicU64::new(0)
|
|
)
|
|
.unwrap_err()
|
|
.contains("mtp.safetensors is not downloaded")
|
|
);
|
|
fs::hard_link(source.join("mtp.safetensors"), root.join("mtp.safetensors")).unwrap();
|
|
let verified = AtomicU64::new(0);
|
|
assert_eq!(
|
|
validate_qwen(
|
|
ManagedArtifactId::Qwen38FlashNext,
|
|
&models,
|
|
&AtomicBool::new(false),
|
|
&verified
|
|
)
|
|
.unwrap(),
|
|
DownloadOutcome::Complete
|
|
);
|
|
assert_eq!(verified.load(Ordering::Relaxed), qwen::total_bytes());
|
|
assert!(qwen::is_installed(&root));
|
|
assert!(crate::engine::validate_qwen_artifacts(&root).is_ok());
|
|
delete_qwen(ManagedArtifactId::Qwen38FlashNext, &models).unwrap();
|
|
assert!(!root.exists());
|
|
fs::remove_dir_all(models).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn issue100_catalog_keeps_vision_and_dspark_checkpoint_specific() {
|
|
let model = ModelChoice::DeepSeekV4FlashVision;
|
|
assert_eq!(ModelChoice::from_id(model.id()), Some(model));
|
|
assert_eq!(
|
|
serde_json::from_str::<ModelChoice>(&serde_json::to_string(&model).unwrap()).unwrap(),
|
|
model
|
|
);
|
|
assert!(model.supports_vision() && model.supports_dspark());
|
|
assert!(!model.supports_integrated_mtp());
|
|
assert_eq!(model.artifacts(true).count(), 3);
|
|
assert_eq!(model.artifacts(false).count(), 2);
|
|
assert_ne!(
|
|
model.dspark_artifact().unwrap().sha256,
|
|
FLASH_0731_DSPARK.sha256
|
|
);
|
|
let defaults = crate::config::ModelPreferences::defaults_for(model);
|
|
assert!(defaults.speculative.dspark_enabled);
|
|
assert!(defaults.speculative.validate(model).is_ok());
|
|
let root = std::env::temp_dir().join(format!("issue100-catalog-{}", std::process::id()));
|
|
fs::create_dir_all(root.join(model.id())).unwrap();
|
|
// Sparse files exercise catalog and marker bookkeeping only. Real GGUF
|
|
// content and hashes are checked by the separate managed-intake test.
|
|
for artifact in [&FLASH_VISION, &FLASH_VISION_DSPARK] {
|
|
File::create(artifact.path(model, &root))
|
|
.unwrap()
|
|
.set_len(artifact.size)
|
|
.unwrap();
|
|
mark_verified(model, artifact, &root).unwrap();
|
|
}
|
|
assert!(!installed_models(&root).contains(&model));
|
|
let missing = engine_artifacts(model, true, &root);
|
|
assert!(
|
|
validate_engine_artifacts(model, true, &missing)
|
|
.unwrap_err()
|
|
.contains("requires its vision encoder")
|
|
);
|
|
File::create(FLASH_VISION_ENCODER.path(model, &root))
|
|
.unwrap()
|
|
.set_len(FLASH_VISION_ENCODER.size)
|
|
.unwrap();
|
|
mark_verified(model, &FLASH_VISION_ENCODER, &root).unwrap();
|
|
assert_eq!(installed_models(&root), [model]);
|
|
let mut artifacts = engine_artifacts(model, true, &root);
|
|
validate_engine_artifacts(model, true, &artifacts).unwrap();
|
|
artifacts.support = Some(FLASH_0731_DSPARK.path(ModelChoice::DeepSeekV4Flash0731, &root));
|
|
assert!(
|
|
validate_engine_artifacts(model, true, &artifacts)
|
|
.unwrap_err()
|
|
.contains("not the expected")
|
|
);
|
|
artifacts = engine_artifacts(model, false, &root);
|
|
artifacts.vision = Some(GLM53_FLASH_VISION.path(ModelChoice::Glm53Flash, &root));
|
|
assert!(
|
|
validate_engine_artifacts(model, false, &artifacts)
|
|
.unwrap_err()
|
|
.contains("not the expected")
|
|
);
|
|
fs::remove_dir_all(root).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "requires all three pinned Vision-Exp GGUF artifacts; performs full SHA256 and intake validation"]
|
|
fn issue100_managed_intake() {
|
|
let source = Path::new("local-eval-results/issue100-artifacts");
|
|
let root = Path::new("local-eval-results/issue100-managed");
|
|
let model = ModelChoice::DeepSeekV4FlashVision;
|
|
fs::create_dir_all(root.join(model.id())).unwrap();
|
|
for id in [
|
|
ManagedArtifactId::DeepSeekV4FlashVisionEncoder,
|
|
ManagedArtifactId::DeepSeekV4FlashVisionDspark,
|
|
ManagedArtifactId::DeepSeekV4FlashVision,
|
|
] {
|
|
let artifact = id.artifact();
|
|
let destination = artifact.path(model, root);
|
|
if !destination.exists() {
|
|
fs::hard_link(source.join(artifact.file_name), &destination).unwrap();
|
|
}
|
|
let verified = AtomicU64::new(0);
|
|
assert_eq!(
|
|
validate_managed_artifact(id, root, &AtomicBool::new(false), &verified).unwrap(),
|
|
DownloadOutcome::Complete
|
|
);
|
|
assert_eq!(verified.load(Ordering::Relaxed), artifact.size);
|
|
eprintln!("verified {} {} bytes", artifact.file_name, artifact.size);
|
|
}
|
|
assert_eq!(installed_models(root), [model]);
|
|
for dspark in [false, true] {
|
|
validate_engine_artifacts(model, dspark, &engine_artifacts(model, dspark, root))
|
|
.unwrap();
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn verification_reports_bytes_read() {
|
|
let path = std::env::temp_dir().join(format!(
|
|
"ds4-server-verify-{}",
|
|
SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_nanos()
|
|
));
|
|
fs::write(&path, b"abc").unwrap();
|
|
let artifact = Artifact {
|
|
label: "test model",
|
|
file_name: "unused",
|
|
repository: "unused",
|
|
size: 3,
|
|
sha256: "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad",
|
|
support: None,
|
|
vision: false,
|
|
};
|
|
let verified_bytes = AtomicU64::new(999);
|
|
|
|
assert_eq!(
|
|
verify(
|
|
&path,
|
|
&artifact,
|
|
ModelChoice::DeepSeekV4Flash0731,
|
|
&AtomicBool::new(false),
|
|
&verified_bytes,
|
|
)
|
|
.unwrap(),
|
|
DownloadOutcome::Complete
|
|
);
|
|
assert_eq!(verified_bytes.load(Ordering::Relaxed), 3);
|
|
fs::remove_file(path).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn managed_artifact_inventory_and_delete_include_partial_files() {
|
|
let models_path = std::env::temp_dir().join(format!(
|
|
"ds4-server-manager-{}",
|
|
SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_nanos()
|
|
));
|
|
let id = ManagedArtifactId::DeepSeekV4Flash0731;
|
|
let partial = id.artifact().partial_path(id.model(), &models_path);
|
|
fs::create_dir_all(partial.parent().unwrap()).unwrap();
|
|
fs::write(&partial, b"part").unwrap();
|
|
|
|
let managed = managed_artifacts(&models_path)
|
|
.into_iter()
|
|
.find(|artifact| artifact.id == id)
|
|
.unwrap();
|
|
assert_eq!(managed.stored, 4);
|
|
assert_eq!(managed.state, ManagedArtifactState::Partial);
|
|
|
|
delete_managed_artifact(id, &models_path).unwrap();
|
|
assert!(!partial.exists());
|
|
fs::remove_dir_all(models_path).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn restart_resumes_at_the_existing_partial_byte() {
|
|
let content = b"restart-resume works";
|
|
let offset = 8;
|
|
let directory = std::env::temp_dir().join(format!(
|
|
"ds4-server-resume-{}",
|
|
SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_nanos()
|
|
));
|
|
fs::create_dir_all(&directory).unwrap();
|
|
let partial = directory.join("model.gguf.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; 2048];
|
|
let count = connection.read(&mut request).unwrap();
|
|
let request = String::from_utf8_lossy(&request[..count]).to_ascii_lowercase();
|
|
assert!(request.contains("range: bytes=8-"));
|
|
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();
|
|
});
|
|
|
|
let outcome = download_url_to_partial_with_timeout(
|
|
&format!("http://{address}/model.gguf"),
|
|
&partial,
|
|
&AtomicBool::new(false),
|
|
false,
|
|
DOWNLOAD_STALL_TIMEOUT,
|
|
)
|
|
.unwrap();
|
|
server.join().unwrap();
|
|
assert_eq!(outcome, DownloadOutcome::Complete);
|
|
assert_eq!(fs::read(&partial).unwrap(), content);
|
|
fs::remove_dir_all(directory).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn stalled_download_times_out_and_keeps_the_partial_file() {
|
|
let directory = std::env::temp_dir().join(format!(
|
|
"ds4-server-stall-{}",
|
|
SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_nanos()
|
|
));
|
|
fs::create_dir_all(&directory).unwrap();
|
|
let partial = directory.join("model.gguf.part");
|
|
fs::write(&partial, b"part").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; 2048];
|
|
let _count = connection.read(&mut request).unwrap();
|
|
connection
|
|
.write_all(
|
|
b"HTTP/1.1 206 Partial Content\r\nContent-Length: 10\r\nContent-Range: bytes 4-13/14\r\nConnection: close\r\n\r\n",
|
|
)
|
|
.unwrap();
|
|
thread::sleep(Duration::from_millis(200));
|
|
});
|
|
|
|
let error = download_url_to_partial_with_timeout(
|
|
&format!("http://{address}/model.gguf"),
|
|
&partial,
|
|
&AtomicBool::new(false),
|
|
false,
|
|
Duration::from_millis(50),
|
|
)
|
|
.unwrap_err();
|
|
server.join().unwrap();
|
|
|
|
assert!(error.contains("timeout"), "{error}");
|
|
assert_eq!(fs::read(&partial).unwrap(), b"part");
|
|
fs::remove_dir_all(directory).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn cancellation_keeps_the_partial_file_for_the_next_run() {
|
|
let directory = std::env::temp_dir().join(format!(
|
|
"ds4-server-cancel-{}",
|
|
SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_nanos()
|
|
));
|
|
let artifact = Artifact {
|
|
label: "test model",
|
|
file_name: "model.gguf",
|
|
repository: "unused",
|
|
size: 10,
|
|
sha256: "unused",
|
|
support: None,
|
|
vision: false,
|
|
};
|
|
let partial = artifact.partial_path(ModelChoice::DeepSeekV4Flash0731, &directory);
|
|
fs::create_dir_all(partial.parent().unwrap()).unwrap();
|
|
fs::write(&partial, b"part").unwrap();
|
|
let cancel = AtomicBool::new(true);
|
|
|
|
assert_eq!(
|
|
download_artifact_with_cancel(
|
|
ModelChoice::DeepSeekV4Flash0731,
|
|
&artifact,
|
|
&directory,
|
|
&cancel,
|
|
&AtomicU64::new(0),
|
|
)
|
|
.unwrap(),
|
|
DownloadOutcome::Stopped
|
|
);
|
|
assert_eq!(fs::read(&partial).unwrap(), b"part");
|
|
fs::remove_dir_all(directory).unwrap();
|
|
}
|
|
}
|