stprobe 0.2.1

A minimal CLI for inspecting safetensors headers
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93

mod common;

use wiremock::matchers::{header, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};

#[test]
fn inspects_remote_files_via_range_requests() {
    let runtime = tokio::runtime::Runtime::new().expect("tokio runtime");
    let server = runtime.block_on(MockServer::start());
    let bytes = common::sample_safetensors_bytes();
    let total_size = bytes.len();
    let (prefix, header_bytes) = common::split_header(&bytes);
    let header_end = 8 + header_bytes.len() - 1;
    let header_range = format!("bytes=8-{header_end}");
    let base_url = server.uri();
    let resolve_path = "/resolve/main/sample.safetensors";
    let cdn_path = "/cdn/sample.safetensors";

    runtime.block_on(async {
        Mock::given(method("GET"))
            .and(path(resolve_path))
            .respond_with(
                ResponseTemplate::new(302)
                    .append_header("Location", format!("{base_url}{cdn_path}")),
            )
            .expect(2)
            .mount(&server)
            .await;

        Mock::given(method("GET"))
            .and(path(cdn_path))
            .and(header("range", "bytes=0-7"))
            .respond_with(
                ResponseTemplate::new(206)
                    .append_header("Content-Range", format!("bytes 0-7/{total_size}"))
                    .set_body_bytes(prefix),
            )
            .expect(1)
            .mount(&server)
            .await;

        Mock::given(method("GET"))
            .and(path(cdn_path))
            .and(header("range", header_range.as_str()))
            .respond_with(
                ResponseTemplate::new(206)
                    .append_header(
                        "Content-Range",
                        format!("bytes 8-{header_end}/{total_size}"),
                    )
                    .set_body_bytes(header_bytes),
            )
            .expect(1)
            .mount(&server)
            .await;
    });

    let url = format!("{base_url}{resolve_path}");
    let report = stprobe::inspect_input(&url).expect("inspect remote safetensors");
    let output = stprobe::render_report(&report);

    assert!(output.contains(&format!("File: {url}")));
    assert!(output.contains("Tensors: 2"));
    assert!(output.contains("Parameters: 4"));
    assert!(output.contains("Tensor-Bytes: 24"));
    assert!(output.contains("  format = pt"));
    assert!(output.contains("  embedding.ids"));
    assert!(output.contains("  embedding.weight"));
}

#[test]
fn reports_servers_without_range_support() {
    let runtime = tokio::runtime::Runtime::new().expect("tokio runtime");
    let server = runtime.block_on(MockServer::start());

    runtime.block_on(async {
        Mock::given(method("GET"))
            .and(path("/sample.safetensors"))
            .and(header("range", "bytes=0-7"))
            .respond_with(ResponseTemplate::new(200).set_body_bytes(vec![0; 8]))
            .expect(1)
            .mount(&server)
            .await;
    });

    let url = format!("{}/sample.safetensors", server.uri());
    let error = stprobe::inspect_input(&url).expect_err("range support error");

    assert_eq!(
        error.to_string(),
        format!("remote server does not support byte range requests: {url}")
    );
}