tensorfs 0.3.58__tar.gz → 0.3.59__tar.gz

{tensorfs-0.3.58 → tensorfs-0.3.59}/Cargo.lock

@@ -439,7 +439,7 @@ checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
 
 [[package]]
 name = "tensorfs-core"
-version = "0.3.58"
+version = "0.3.59"
 dependencies = [
  "fs2",
  "libc",
@@ -454,7 +454,7 @@ dependencies = [
 
 [[package]]
 name = "tensorfs-py"
-version = "0.3.58"
+version = "0.3.59"
 dependencies = [
  "pyo3",
  "tensorfs-core",

{tensorfs-0.3.58 → tensorfs-0.3.59}/Cargo.toml

@@ -3,5 +3,5 @@ resolver = "2"
 members = ["crates/tensorfs-core", "crates/tensorfs-py"]
 
 [workspace.package]
-version = "0.3.58"
+version = "0.3.59"
 edition = "2021"

{tensorfs-0.3.58 → tensorfs-0.3.59}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tensorfs
-Version: 0.3.58
+Version: 0.3.59
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Rust

{tensorfs-0.3.58 → tensorfs-0.3.59}/crates/tensorfs-core/src/catalog.rs

@@ -911,9 +911,14 @@ impl Catalog {
         object: &crate::ids::ObjectRef,
         publish: impl FnOnce() -> Result<Option<crate::store::VRecord>>,
     ) -> Result<bool> {
+        // One admission at a time per catalog in this process, holding the write lock from
+        // before publication: streams landing objects together queue here instead of in
+        // SQLite's busy timeout, where a lost race refused bytes already verified.
+        let turn = admissions(&self.path);
+        let _turn = turn.lock().unwrap_or_else(|error| error.into_inner());
         let mut connection = self.connection()?;
         let transaction = connection
-            .transaction()
+            .transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)
             .map_err(|error| sql("begin blob admission", error))?;
         if let Some(operation) = operation {
             hold_on(
@@ -1079,6 +1084,17 @@ impl Catalog {
     }
 }
 
+fn admissions(path: &Path) -> Arc<Mutex<()>> {
+    static TURNS: OnceLock<Mutex<HashMap<PathBuf, Arc<Mutex<()>>>>> = OnceLock::new();
+    TURNS
+        .get_or_init(Default::default)
+        .lock()
+        .unwrap_or_else(|error| error.into_inner())
+        .entry(path.to_path_buf())
+        .or_default()
+        .clone()
+}
+
 fn hold_on(connection: &Connection, operation: &str, key: &str, length: u64) -> Result<()> {
     let kind = if key.starts_with("blobs/") {
         "blob"

{tensorfs-0.3.58 → tensorfs-0.3.59}/crates/tensorfs-core/src/transport/http.rs

@@ -60,6 +60,11 @@ pub enum Judge {
     /// which is the owner's rule read literally: the lack of measured progress, never a
     /// number of seconds.
     Transfer,
+    /// An UPLOAD. Every wait — connect, handshake, each write, and the answer after the
+    /// last body byte — is judged by this exchange's own silence against the ledger's
+    /// patience. Nothing it moves lands here, so it neither counts as progress nor teaches
+    /// the pace.
+    Silence,
     /// Neither. The caller's deadline is the only bound, which is what a foreign-source
     /// read with its own per-part retry loop above it already had.
     Deadline,
@@ -133,7 +138,7 @@ fn transfer_verdict(
 /// progress may span arbitrarily many windows without exhausting a wall budget.
 fn stream_verdict(judge: Judge, ledger: &Ledger, advanced: Instant, phase: &str) -> Result<()> {
     ledger.check_running()?;
-    if judge != Judge::Stream || advanced.elapsed() <= ledger.patience() {
+    if !judge.silent_is_stall() || advanced.elapsed() <= ledger.patience() {
         return Ok(());
     }
     refuse(
@@ -147,6 +152,12 @@ fn stream_verdict(judge: Judge, ledger: &Ledger, advanced: Instant, phase: &str)
     )
 }
 
+impl Judge {
+    fn silent_is_stall(self) -> bool {
+        matches!(self, Judge::Stream | Judge::Silence)
+    }
+}
+
 pub struct Client {
     tls: Arc<rustls::ClientConfig>,
     idle: IdlePool,
@@ -712,7 +723,7 @@ impl BodyReader<'_> {
                 Err(e) if would_block(&e) => {
                     transfer_verdict(self.judge, self.ledger, self.since, "a response body")?;
                     let silent = self.advanced.elapsed();
-                    if self.judge == Judge::Stream && silent > self.ledger.patience() {
+                    if self.judge.silent_is_stall() && silent > self.ledger.patience() {
                         self.stalled = true;
                         let _ = self
                             .wire

tensorfs-0.3.59/crates/tensorfs-core/src/transport/memory.rs

@@ -0,0 +1,92 @@
+//! How many transfer streams fit in the memory this process can actually use. There is no
+//! assumed host: a 4 GB pod runs the same pull, source download or upload window with fewer
+//! streams, never with N × buffers it does not have.
+use std::fs;
+use std::path::{Path, PathBuf};
+
+/// The most memory one transfer stream holds, whatever its object's size: one 1 MiB body
+/// chunk (the read, write, hash or upload buffer), the 64 KiB framing buffer, and rustls's
+/// bounded record buffers, rounded up to a whole 2 MiB.
+pub const STREAM_MEMORY: u64 = 2 << 20;
+
+/// Bytes this process can still use now: the tightest of its cgroup's limit less the usage
+/// that cannot be reclaimed (v2 at every ancestor, else v1) and the host's `MemAvailable`.
+/// `None` when none of them is readable.
+pub fn memory_available() -> Option<u64> {
+    [cgroup_v2(), cgroup_v1(), mem_available()]
+        .into_iter()
+        .flatten()
+        .min()
+}
+
+/// Concurrent transfer streams the memory available now allows: at most `wanted`, at
+/// least one.
+pub fn transfer_streams(wanted: usize) -> usize {
+    streams_within(wanted, memory_available())
+}
+
+/// `transfer_streams` against a stated budget.
+pub fn streams_within(wanted: usize, available: Option<u64>) -> usize {
+    let fits = available.map_or(usize::MAX, |bytes| {
+        usize::try_from(bytes / STREAM_MEMORY).unwrap_or(usize::MAX)
+    });
+    wanted.min(fits).max(1)
+}
+
+fn number(path: &Path) -> Option<u64> {
+    fs::read_to_string(path).ok()?.trim().parse().ok()
+}
+
+/// `limit − (usage − page cache)`: page cache is reclaimed under pressure, so bytes a
+/// download has written do not count against the next stream.
+fn headroom(dir: &Path, limit: &str, usage: &str, cache: &str) -> Option<u64> {
+    let limit = number(&dir.join(limit))?;
+    let usage = number(&dir.join(usage))?;
+    let cached = fs::read_to_string(dir.join("memory.stat"))
+        .ok()
+        .and_then(|stat| {
+            stat.lines()
+                .find_map(|line| line.strip_prefix(cache)?.strip_prefix(' ')?.parse().ok())
+        })
+        .unwrap_or(0);
+    Some(limit.saturating_sub(usage.saturating_sub(cached)))
+}
+
+fn cgroup_v2() -> Option<u64> {
+    let root = Path::new("/sys/fs/cgroup");
+    let own = fs::read_to_string("/proc/self/cgroup").ok()?;
+    let own = own.lines().find_map(|line| line.strip_prefix("0::"))?;
+    let mut dir = root.join(own.trim().trim_start_matches('/'));
+    let mut tightest: Option<u64> = None;
+    loop {
+        if let Some(left) = headroom(&dir, "memory.max", "memory.current", "file") {
+            tightest = Some(tightest.map_or(left, |held| held.min(left)));
+        }
+        if dir == root || !dir.pop() || !dir.starts_with(root) {
+            return tightest;
+        }
+    }
+}
+
+fn cgroup_v1() -> Option<u64> {
+    let dir = PathBuf::from("/sys/fs/cgroup/memory");
+    headroom(
+        &dir,
+        "memory.limit_in_bytes",
+        "memory.usage_in_bytes",
+        "total_cache",
+    )
+}
+
+fn mem_available() -> Option<u64> {
+    let text = fs::read_to_string("/proc/meminfo").ok()?;
+    let kib: u64 = text
+        .lines()
+        .find_map(|line| line.strip_prefix("MemAvailable:"))?
+        .trim()
+        .strip_suffix("kB")?
+        .trim()
+        .parse()
+        .ok()?;
+    kib.checked_mul(1024)
+}

{tensorfs-0.3.58 → tensorfs-0.3.59}/crates/tensorfs-core/src/transport/mod.rs

@@ -35,6 +35,7 @@ mod api;
 mod http;
 mod hub;
 mod ledger;
+mod memory;
 mod policy;
 pub(crate) mod prefix;
 mod pull;
@@ -55,6 +56,7 @@ pub use ledger::{
     Deadline, Ledger, PullCancellation, FETCH_ATTEMPTS, SAMPLE_SECONDS, STALL_FACTOR,
     STILL_SAMPLES, STREAMS,
 };
+pub use memory::{memory_available, streams_within, transfer_streams, STREAM_MEMORY};
 pub use policy::{base_host, blocked, is_loopback_origin, resolve, CheckedUrl, SourcePolicy};
 pub use pull::{
     fetch_authorized, fetch_url, fetch_wanted, pull, FetchObserver, Fetched, ObjectSource,

{tensorfs-0.3.58 → tensorfs-0.3.59}/crates/tensorfs-core/src/transport/prefix.rs

@@ -479,7 +479,7 @@ pub(crate) fn fetch_sources(
     };
     download.report(0, held);
     std::thread::scope(|scope| {
-        for _ in 0..options.streams.clamp(1, 128) {
+        for _ in 0..super::transfer_streams(options.streams.clamp(1, 128)) {
             scope.spawn(|| download.work());
         }
     });

{tensorfs-0.3.58 → tensorfs-0.3.59}/crates/tensorfs-core/src/transport/pull.rs

@@ -980,7 +980,7 @@ pub(super) fn fetch_wanted_bounded(
     // invariant it proved — the grant set covers the wanted set exactly — is enforced per
     // mint instead: presign authorizes every digest it was asked or refuses, never a silent
     // subset, and a supplier with no URL for an object refuses in the same words.
-    let streams = streams.clamp(1, 128);
+    let streams = super::memory::transfer_streams(streams.clamp(1, 128));
     let recovery = super::speculate::Recovery::new(store, plan, streams, download_budget);
     // Recovery reserves an object connection, not an authorization worker.
     let workers = streams.min(plan.wanted.len().max(1));

{tensorfs-0.3.58 → tensorfs-0.3.59}/crates/tensorfs-core/src/transport/push.rs

@@ -10,11 +10,12 @@
 //! successful end state for a content-addressed write, not a failure.
 //!
 //! The object is RE-SENT when the socket goes silent or the store answers a retryable
-//! status, `PUT_ATTEMPTS` times, exactly as the pull re-asks for one object. The write
-//! silence bound is the ledger's floor — the same derived noise floor `egress.py` used —
-//! never an invented constant. A retryable answer is re-sent after the wait the store
-//! states (`Retry-After`), which spends no attempt, or after the ledger's patience, which
-//! does. Authority stays where it was: `podweights` (or whoever
+//! status, `PUT_ATTEMPTS` times, exactly as the pull re-asks for one object. Every wait of
+//! the exchange — the answer after the last body byte included — is judged by its own
+//! silence against the ledger's patience (`Judge::Silence`), and one write that moves
+//! nothing by the ledger's floor; never an invented constant. A retryable answer is re-sent
+//! after the wait the store states (`Retry-After`), which spends no attempt, or after the
+//! ledger's patience, which does. Authority stays where it was: `podweights` (or whoever
 //! minted the grant) decided WHETHER these bytes may go; this module only moves them.
 //!
 //! **A grant is a URL AND the headers its signature covers, together.** tensorhub signs
@@ -228,10 +229,9 @@ pub fn push_object(
             Some(payload),
             deadline,
             ledger,
-            // The push plane keeps its own derived write bound below and gains no plane
-            // here: it serves no hub call, so it has no transfer whose stillness could
-            // judge it.
-            http::Judge::Deadline,
+            // A peer that stops answering — during the handshake, a write, or after the
+            // last body byte — is judged by that silence, never waited on forever.
+            http::Judge::Silence,
             Some(ledger.floor()),
             Code::TRANSFER_FAILED,
         );

{tensorfs-0.3.58 → tensorfs-0.3.59}/crates/tensorfs-core/tests/source_download.rs

@@ -19,8 +19,8 @@ use tensorfs_core::{
     providers::{self, Provenance, Pulled, Resolution, ResolvedMember},
     store::{Fault, Store},
     transport::{
-        push_object, Anonymous, Deadline, Ledger, PullCancellation, SourceDownload, SourcePolicy,
-        UploadGrant,
+        memory_available, push_object, streams_within, transfer_streams, Anonymous, Deadline,
+        Ledger, PullCancellation, SourceDownload, SourcePolicy, UploadGrant, STREAM_MEMORY,
     },
 };
 
@@ -55,6 +55,8 @@ enum Act {
     Status(u16, Vec<(String, String)>),
     /// Send the head and this many body bytes, then drop the connection.
     Cut(usize),
+    /// Read a PUT body, then say nothing, ever.
+    Silent,
 }
 
 struct Request {
@@ -177,13 +179,19 @@ impl Origin {
             let number = self.asks.fetch_add(1, Ordering::Relaxed) + 1;
             let act = (self.decide)(&request, number);
             if request.method == "PUT" {
-                let length: usize = header("content-length").unwrap().parse().unwrap();
-                let mut sink = vec![0u8; length];
-                if socket.read_exact(&mut sink).is_err() {
+                let length: u64 = header("content-length").unwrap().parse().unwrap();
+                // Discarded as it arrives: this origin's memory is not the client's.
+                let taken = std::io::copy(&mut (&mut socket).take(length), &mut std::io::sink());
+                if taken.ok() != Some(length) {
                     return;
                 }
                 self.puts.fetch_add(1, Ordering::Relaxed);
                 match act {
+                    Act::Silent => {
+                        // Hold the socket until the client gives up on it.
+                        let _ = socket.read(&mut [0u8; 1]);
+                        return;
+                    }
                     Act::Status(status, extra) => {
                         if !reply(&mut socket, status, &extra, b"retry") {
                             return;
@@ -684,6 +692,45 @@ fn pushed_store(name: &str) -> (std::path::PathBuf, Store, ObjectRef) {
     (dir, store, object)
 }
 
+#[test]
+fn a_push_whose_store_goes_silent_after_the_body_is_judged_not_waited_on() {
+    let (dir, store, object) = pushed_store("push-silent");
+    let origin = Origin::start(vec![], 0, |request, _| {
+        if request.method == "PUT" {
+            Act::Silent
+        } else {
+            Act::Serve
+        }
+    });
+    let ledger = Ledger::with_resolution(Duration::from_millis(20));
+    let started = Instant::now();
+    let refusal = push_object(
+        &store,
+        &object.sha256,
+        false,
+        &UploadGrant::to(&format!("{}/put/object", origin.base)),
+        &policy(),
+        &Anonymous,
+        Deadline::none(),
+        &ledger,
+        "",
+    )
+    .unwrap_err();
+    assert_eq!(refusal.code, Code::TRANSFER_FAILED);
+    assert!(refusal.detail.contains("response head"), "{refusal}");
+    assert_eq!(
+        origin.puts.load(Ordering::Relaxed),
+        4,
+        "each attempt sent the whole body"
+    );
+    assert!(
+        started.elapsed() < Duration::from_secs(30),
+        "{:?}",
+        started.elapsed()
+    );
+    fs::remove_dir_all(dir).unwrap();
+}
+
 #[test]
 fn a_push_told_to_wait_waits_and_then_lands() {
     let (dir, store, object) = pushed_store("push-wait");
@@ -713,6 +760,210 @@ fn a_push_told_to_wait_waits_and_then_lands() {
     fs::remove_dir_all(dir).unwrap();
 }
 
+#[test]
+fn concurrent_admissions_queue_instead_of_refusing() {
+    let dir = root("admissions");
+    let store = Store::init(&dir).unwrap();
+    thread::scope(|scope| {
+        for worker in 0..16u64 {
+            let store = &store;
+            scope.spawn(move || {
+                for index in 0..4 {
+                    let bytes = body(1000 + worker * 4 + index, 64 << 10);
+                    let object = ObjectRef::of(&bytes);
+                    store
+                        .put_stream(&mut bytes.as_slice(), Some(&object), &Fault::default())
+                        .unwrap();
+                    assert!(store.record_valid(&object.sha256).is_ok());
+                }
+            });
+        }
+    });
+    fs::remove_dir_all(dir).unwrap();
+}
+
+/// Peak resident memory since `/proc/self/clear_refs` was last reset, in bytes.
+fn peak_resident() -> u64 {
+    let status = fs::read_to_string("/proc/self/status").unwrap();
+    let kib: u64 = status
+        .lines()
+        .find_map(|line| line.strip_prefix("VmHWM:"))
+        .unwrap()
+        .trim()
+        .trim_end_matches("kB")
+        .trim()
+        .parse()
+        .unwrap();
+    kib * 1024
+}
+
+/// Run one `*_child` test of this binary in its own process, optionally inside `scope`.
+fn run_child(test: &str, dir: &std::path::Path, scope: &[&str]) -> std::process::Output {
+    let exe = std::env::current_exe().unwrap();
+    let mut command = match scope.split_first() {
+        Some((program, args)) => {
+            let mut command = Command::new(program);
+            command.args(args).arg(&exe);
+            command
+        }
+        None => Command::new(&exe),
+    };
+    command
+        .args(["--exact", test, "--nocapture"])
+        .env("TFS_DOWNLOAD_CHILD", dir)
+        .output()
+        .unwrap()
+}
+
+#[test]
+fn push_memory_child() {
+    let Some(dir) = std::env::var_os("TFS_DOWNLOAD_CHILD") else {
+        return;
+    };
+    let store = Store::init(std::path::Path::new(&dir)).unwrap();
+    let objects: Vec<ObjectRef> = (0..16u64)
+        .map(|seed| {
+            let bytes = body(300 + seed, 16 * MIB);
+            let object = ObjectRef::of(&bytes);
+            store
+                .put_stream(&mut bytes.as_slice(), Some(&object), &Fault::default())
+                .unwrap();
+            object
+        })
+        .collect();
+    let origin = Origin::start(vec![], 0, |_, _| Act::Serve);
+    // Forget the objects' construction; measure only the pushes.
+    fs::write("/proc/self/clear_refs", "5").unwrap();
+    let before = peak_resident();
+    thread::scope(|scope| {
+        for object in &objects {
+            let (store, origin) = (&store, &origin);
+            scope.spawn(move || {
+                push_object(
+                    store,
+                    &object.sha256,
+                    false,
+                    &UploadGrant::to(&format!("{}/put/{}", origin.base, object.sha256)),
+                    &policy(),
+                    &Anonymous,
+                    Deadline::none(),
+                    &Ledger::with_resolution(Duration::from_millis(200)),
+                    "",
+                )
+                .unwrap();
+            });
+        }
+    });
+    let grew = peak_resident().saturating_sub(before);
+    println!("PUSHED 16 x 16 MiB, peak resident grew {grew} B");
+    assert_eq!(origin.puts.load(Ordering::Relaxed), 16);
+    // 256 MiB crossed the sockets; the pushes may hold their streams' buffers and no more.
+    assert!(
+        grew < 16 * STREAM_MEMORY + (16 << 20),
+        "pushes buffered {grew} B"
+    );
+}
+
+#[test]
+fn sixteen_concurrent_pushes_hold_stream_buffers_not_objects() {
+    let dir = root("push-memory");
+    let output = run_child("push_memory_child", &dir, &[]);
+    let text = String::from_utf8_lossy(&output.stdout).to_string()
+        + &String::from_utf8_lossy(&output.stderr);
+    assert!(output.status.success() && text.contains("PUSHED"), "{text}");
+    println!(
+        "{}",
+        text.lines().find(|line| line.contains("PUSHED")).unwrap()
+    );
+    fs::remove_dir_all(dir).unwrap();
+}
+
+#[test]
+fn streams_are_sized_to_the_memory_budget() {
+    assert_eq!(streams_within(16, None), 16);
+    assert_eq!(streams_within(16, Some(1 << 40)), 16);
+    assert_eq!(streams_within(16, Some(5 * STREAM_MEMORY)), 5);
+    assert_eq!(
+        streams_within(16, Some(0)),
+        1,
+        "a starved host still moves one stream"
+    );
+    assert!(memory_available().is_some_and(|bytes| bytes > 0));
+}
+
+#[test]
+fn bounded_memory_child() {
+    let Some(dir) = std::env::var_os("TFS_DOWNLOAD_CHILD") else {
+        return;
+    };
+    let available = memory_available().unwrap();
+    let streams = transfer_streams(1024);
+    println!("BUDGET {available} B allows {streams} streams");
+    // The scope's 64 MiB, not the host's gigabytes, decides.
+    assert!(
+        available < 64 << 20 && streams < 32,
+        "{available} B, {streams} streams"
+    );
+    let name = "model.safetensors";
+    let origin = Origin::start(vec![(name, body(77, 16 * MIB))], 16 << 20, |_, _| {
+        Act::Serve
+    });
+    let store = Store::init(std::path::Path::new(&dir)).unwrap();
+    materialize(
+        &store,
+        &owner("bounded"),
+        &origin.resolution(&[name]),
+        &options(MIB, 1024),
+        &|_, _| {},
+    )
+    .unwrap();
+    verify(&store, &origin, &[name]);
+    let live = origin.most_live.load(Ordering::SeqCst);
+    println!("DOWNLOADED with at most {live} live streams");
+    assert!(
+        live <= streams,
+        "{live} streams against a budget of {streams}"
+    );
+}
+
+#[test]
+fn a_memory_limited_scope_downloads_with_fewer_streams() {
+    let scope = [
+        "systemd-run",
+        "--user",
+        "--scope",
+        "-q",
+        "-p",
+        "MemoryMax=64M",
+    ];
+    let usable = Command::new(scope[0])
+        .args(&scope[1..])
+        .arg("true")
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .status()
+        .is_ok_and(|status| status.success());
+    if !usable {
+        eprintln!("SKIPPED: no systemd user manager to create a memory-limited cgroup");
+        return;
+    }
+    let dir = root("bounded-memory");
+    let output = run_child("bounded_memory_child", &dir, &scope);
+    let text = String::from_utf8_lossy(&output.stdout).to_string()
+        + &String::from_utf8_lossy(&output.stderr);
+    assert!(
+        output.status.success() && text.contains("DOWNLOADED"),
+        "{text}"
+    );
+    for line in text
+        .lines()
+        .filter(|line| line.contains("BUDGET") || line.contains("DOWNLOADED"))
+    {
+        println!("{line}");
+    }
+    let _ = fs::remove_dir_all(dir);
+}
+
 /// Throughput on a paced loopback origin (each connection capped like one CDN stream).
 /// `cargo test --test source_download -- --ignored --nocapture throughput`
 #[test]

{tensorfs-0.3.58 → tensorfs-0.3.59}/crates/tensorfs-py/src/lib.rs

@@ -119,6 +119,8 @@ fn _ext(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(store::fetch_complete_cozytensors, m)?)?;
     m.add_class::<store::PyPullCancellation>()?;
     m.add_function(wrap_pyfunction!(store::pull, m)?)?;
+    m.add_function(wrap_pyfunction!(store::transfer_streams, m)?)?;
+    m.add("STREAM_MEMORY", tensorfs_core::transport::STREAM_MEMORY)?;
     m.add_class::<store::PyRepoObjectCache>()?;
     m.add_class::<store::PyStore>()?;
     m.add_class::<store::PyOperation>()?;