From f824f669498d6d64b1d2a7481e9516964b7f8633 Mon Sep 17 00:00:00 2001 From: EthanRosenthal Date: Sat, 3 Oct 2026 13:56:17 +0000 Subject: [PATCH 1/3] Allocate CPU output buffers outside torch's allocator Decoded frame and sample buffers were allocated with at::empty, i.e. through torch's CPU allocator. On the aarch64 Linux wheels that allocator is mimalloc, which does not return a large block freed by a thread other than the one that allocated it: the pages are MADV_FREE'd but stay mapped and resident until the allocating thread exits. Data-loading pipelines decode on worker threads and drop the tensors elsewhere, so every decode left its whole output buffer (hundreds of MiB to GiBs) in resident memory. Allocate the buffers with posix_memalign instead and hand them to torch with at::from_blob plus a free() deleter, which runs on whichever thread drops the last reference; glibc munmaps large blocks immediately regardless of thread. tch-rs has no from_blob-with-deleter wrapper, so build.rs compiles a small C++ shim (csrc/) against the libtorch headers, located the same way torch-sys locates them. Tensor shape, dtype, strides and device are unchanged; CUDA outputs still go through torch. AVTENSOR_OUTPUT_ALLOCATOR=torch restores the previous behaviour. benchmarks/cross_thread_retention.py reproduces the problem: 12 decodes of a 20 s 720p clip (1.24 GiB each) on a 4-thread pool, freed on the main thread, left 14.9 GiB resident (all LazyFree) with at::empty on an aarch64 torch 2.11 build, and none with this change, at the same decode time. --- CLAUDE.md | 8 +- Cargo.lock | 2 + Cargo.toml | 6 + README.md | 11 + benchmarks/cross_thread_retention.py | 176 ++++++++++++++++ build.rs | 125 +++++++++++ csrc/tensor_from_blob.cpp | 46 ++++ pyproject.toml | 2 +- src/alloc.rs | 300 +++++++++++++++++++++++++++ src/decoder/mod.rs | 4 +- src/lib.rs | 1 + 11 files changed, 677 insertions(+), 4 deletions(-) create mode 100644 benchmarks/cross_thread_retention.py create mode 100644 build.rs create mode 100644 csrc/tensor_from_blob.cpp create mode 100644 src/alloc.rs diff --git a/CLAUDE.md b/CLAUDE.md index 50c280b..fd3b2cb 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -12,7 +12,13 @@ generated locally with FFmpeg; end-to-end cloud-storage tests live in the consuming project. `tch` comes from crates.io; its version maps 1:1 to a libtorch version -(0.24.0 ↔ torch 2.11), so bump them together. +(0.24.0 ↔ torch 2.11), so bump them together, along with the `torch-sys` +pin (same version as `tch`). + +`build.rs` compiles `csrc/tensor_from_blob.cpp` against the libtorch headers +(located the same way torch-sys locates them); `src/alloc.rs` uses it to give +decoded-output buffers to torch via `from_blob` with a deleter instead of +`at::empty`. Keep the two in sync if the C API changes. ## Git diff --git a/Cargo.lock b/Cargo.lock index 49b8abc..5bc185a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -146,6 +146,7 @@ dependencies = [ "aws-runtime", "aws-sdk-s3", "bytes", + "cc", "derive_builder", "env_logger", "google-cloud-gax", @@ -164,6 +165,7 @@ dependencies = [ "test-case", "thiserror 2.0.15", "tokio", + "torch-sys", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 1b36ec1..176014f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,6 +28,8 @@ pyo3 = "0.25.0" pyo3-log = "0.12.4" rsmpeg = { version = "0.17.0", features = ["ffmpeg7_1"] } tch = { version = "0.26.0", features = ["python-extension"] } +# Same version as tch; only used for the opaque `C_tensor` handle in src/alloc.rs. +torch-sys = "0.26.0" tempfile = "3.20.0" thiserror = "2.0.15" libloading = "0.8" @@ -38,6 +40,10 @@ tokio = { version = "1.47.1", features = ["time"] } extension-module = ["pyo3/extension-module"] default = ["extension-module"] +[build-dependencies] +# Compiles csrc/tensor_from_blob.cpp against the libtorch headers. +cc = "1.2" + [dev-dependencies] mel_spec = "0.3.3" ndarray = "0.16.1" diff --git a/README.md b/README.md index bfe75d0..976bdf8 100644 --- a/README.md +++ b/README.md @@ -416,6 +416,15 @@ provider-specific support — the URL is handed to FFmpeg's http protocol. - **`s3://` works with AWS but not with your S3-compatible store** — check whether the store requires path-style addressing (`AVTENSOR_S3_FORCE_PATH_STYLE=1`) and that `AWS_ENDPOINT_URL_S3` is set. +- **Resident memory grows when decoded tensors are freed on a different + thread than the one that decoded them** — torch's CPU allocator is not + always plain malloc (the aarch64 Linux wheels embed mimalloc, which keeps + blocks freed cross-thread mapped and resident until the allocating thread + exits). avtensor therefore allocates its output buffers with the system + allocator and hands them to torch via `from_blob` with a deleter, so a free + is an immediate `munmap` on any thread. `AVTENSOR_OUTPUT_ALLOCATOR=torch` + switches back to `torch.empty`-style allocation; + `benchmarks/cross_thread_retention.py` measures the difference. ## Contributing @@ -423,9 +432,11 @@ Issues and pull requests are welcome. Project layout: ``` src/ + alloc.rs output-buffer allocation (system allocator + torch from_blob) decoder/ demuxing, decoding, filter graphs (mod.rs), cloud AVIO reader (io.rs) ffi/ PyO3 bindings: request/response types, decode_asset util/ gcs/s3 URI handling, memory, test media generation +csrc/ C++ shim for the one libtorch call tch-rs does not wrap (built by build.rs) avtensor.pyi Python type stubs, shipped with the wheel ``` diff --git a/benchmarks/cross_thread_retention.py b/benchmarks/cross_thread_retention.py new file mode 100644 index 0000000..8decd9e --- /dev/null +++ b/benchmarks/cross_thread_retention.py @@ -0,0 +1,176 @@ +"""Measure memory retained after decoded tensors are freed on another thread. + +Decodes one clip repeatedly, first allocating and dropping on the main thread, +then decoding on worker threads and dropping on the main thread (the shape of a +typical data-loading pipeline), and reports the process's resident memory after +each phase from ``/proc/self/smaps_rollup``. ``LazyFree`` counts pages that were +``MADV_FREE``'d but are still resident, which is how a thread-caching allocator +keeps a cross-thread free alive. + +Run it twice to compare the output allocators:: + + python benchmarks/cross_thread_retention.py --clip clip.mp4 + AVTENSOR_OUTPUT_ALLOCATOR=torch python benchmarks/cross_thread_retention.py --clip clip.mp4 + +Without ``--clip`` a synthetic 1280x720 H.264 clip with audio is generated with +``ffmpeg`` (``$FFMPEG`` or the one on ``PATH``). Linux only (reads ``/proc``). +""" + +from __future__ import annotations + +import argparse +import gc +import os +import shutil +import subprocess +import sys +import tempfile +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import avtensor +import torch + +GIB = 2**30 + + +def smaps_rollup() -> dict[str, float]: + """Rss / Anonymous / LazyFree of this process in GiB.""" + wanted = {"Rss", "Anonymous", "LazyFree"} + out: dict[str, float] = dict.fromkeys(wanted, 0.0) + with open("/proc/self/smaps_rollup") as f: + for line in f: + key, _, rest = line.partition(":") + if key in wanted: + out[key] = int(rest.split()[0]) * 1024 / GIB + return out + + +def generate_clip(path: Path, seconds: float, size: str, fps: int) -> None: + ffmpeg = os.environ.get("FFMPEG") or shutil.which("ffmpeg") + if ffmpeg is None: + sys.exit("ffmpeg not found (set FFMPEG or PATH), or pass --clip instead") + cmd = [ + ffmpeg, + "-hide_banner", + "-loglevel", + "error", + "-y", + "-f", + "lavfi", + "-i", + f"testsrc2=size={size}:rate={fps}:duration={seconds}", + "-f", + "lavfi", + "-i", + f"sine=frequency=440:sample_rate=48000:duration={seconds}", + "-c:v", + "libx264", + "-preset", + "ultrafast", + "-pix_fmt", + "yuv420p", + "-c:a", + "aac", + str(path), + ] + subprocess.run(cmd, check=True) + + +def decode(clip: str, with_audio: bool) -> list[torch.Tensor]: + """Decode the clip and return its output tensors (so the caller frees them).""" + request = avtensor.MediaDecodeRequest( + clip, + video_stream=avtensor.VideoStreamRequest(), + audio_streams=[avtensor.AudioStreamRequest()] if with_audio else None, + ) + return [s["data"] for s in avtensor.decode_asset(request)] + + +def nbytes(tensors: list[torch.Tensor]) -> int: + return sum(t.numel() * t.element_size() for t in tensors) + + +def report( + tag: str, base: dict[str, float], times: list[float], bytes_out: int +) -> None: + m = smaps_rollup() + mean_ms = 1000 * sum(times) / len(times) if times else float("nan") + print( + f"{tag:<34} rss={m['Rss']:6.2f} anon={m['Anonymous']:6.2f} " + f"lazyfree={m['LazyFree']:6.2f} GiB " + f"(+{m['Rss'] - base['Rss']:.2f} GiB over start) " + f"decode {mean_ms:7.1f} ms output {bytes_out / GIB:.2f} GiB/decode", + flush=True, + ) + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Measure memory retained after cross-thread frees of decoded tensors." + ) + parser.add_argument("--clip", type=Path, help="media file to decode") + parser.add_argument("--decodes", type=int, default=12, help="decodes per phase") + parser.add_argument("--threads", type=int, default=4, help="worker threads") + parser.add_argument( + "--seconds", type=float, default=20.0, help="synthetic clip length" + ) + parser.add_argument("--size", default="1280x720", help="synthetic clip size") + parser.add_argument("--fps", type=int, default=24, help="synthetic clip frame rate") + parser.add_argument( + "--settle", type=float, default=0.3, help="seconds to wait after each free" + ) + args = parser.parse_args() + + allocator = os.environ.get("AVTENSOR_OUTPUT_ALLOCATOR", "system (default)") + print(f"AVTENSOR_OUTPUT_ALLOCATOR={allocator}", flush=True) + + tmpdir: tempfile.TemporaryDirectory[str] | None = None + clip = args.clip + if clip is None: + tmpdir = tempfile.TemporaryDirectory() + clip = Path(tmpdir.name) / "clip.mp4" + generate_clip(clip, args.seconds, args.size, args.fps) + clip = str(clip) + + with_audio = bool(avtensor.probe_asset(clip)["audio_streams"]) + + # Warm up (codec init, allocator arenas) so the baseline is steady. + bytes_out = nbytes(decode(clip, with_audio)) + gc.collect() + base = smaps_rollup() + report("start", base, [], bytes_out) + + times: list[float] = [] + for _ in range(args.decodes): + t0 = time.monotonic() + tensors = decode(clip, with_audio) + times.append(time.monotonic() - t0) + del tensors + time.sleep(args.settle) + report(f"{args.decodes}x decode+free on main thread", base, times, bytes_out) + + pool = ThreadPoolExecutor(max_workers=args.threads) + times = [] + for _ in range(args.decodes): + t0 = time.monotonic() + tensors = pool.submit(decode, clip, with_audio).result() + times.append(time.monotonic() - t0) + del tensors # freed on the main thread, not the decoding thread + time.sleep(args.settle) + report(f"{args.decodes}x decode in pool, free on main", base, times, bytes_out) + + time.sleep(5.0) + report("after 5 s idle", base, times, bytes_out) + + pool.shutdown(wait=True) + time.sleep(0.5) + report("after worker threads exited", base, times, bytes_out) + + if tmpdir is not None: + tmpdir.cleanup() + + +if __name__ == "__main__": + main() diff --git a/build.rs b/build.rs new file mode 100644 index 0000000..e465e1e --- /dev/null +++ b/build.rs @@ -0,0 +1,125 @@ +//! Compiles the C++ shim in `csrc/` that exposes `at::from_blob` with a +//! deleter to Rust (used by `src/alloc.rs`). tch-rs only wraps the non-owning +//! overload, so this is the one piece of libtorch we have to reach through +//! C++. The include directories and C++ ABI are resolved the same way +//! torch-sys resolves them, so both shims link against the same libtorch. + +use std::env; +use std::path::PathBuf; +use std::process::Command; + +const PYTHON_PROBE: &str = r" +import torch +from torch.utils import cpp_extension +print('CXX11_ABI:', int(torch._C._GLIBCXX_USE_CXX11_ABI)) +for path in cpp_extension.include_paths(): + print('INCLUDE:', path) +"; + +struct LibtorchConfig { + include_dirs: Vec, + cxx11_abi: String, +} + +fn env_var_rerun(name: &str) -> Option { + println!("cargo:rerun-if-env-changed={name}"); + env::var(name).ok() +} + +/// Mirrors torch-sys: `python` inside a virtualenv, `python3` otherwise, +/// unless maturin/PyO3 point at a specific interpreter. +fn python_interpreter() -> PathBuf { + if let Some(p) = env_var_rerun("PYO3_PYTHON") { + return PathBuf::from(p); + } + if env::var_os("VIRTUAL_ENV").is_some() { + PathBuf::from("python") + } else { + PathBuf::from("python3") + } +} + +fn config_from_python() -> Result { + let python = python_interpreter(); + let output = Command::new(&python) + .arg("-c") + .arg(PYTHON_PROBE) + .output() + .map_err(|e| format!("running {python:?}: {e}"))?; + if !output.status.success() { + return Err(format!( + "{python:?} failed: {}", + String::from_utf8_lossy(&output.stderr) + )); + } + let stdout = String::from_utf8_lossy(&output.stdout); + let mut include_dirs = vec![]; + let mut cxx11_abi = None; + for line in stdout.lines() { + if let Some(path) = line.strip_prefix("INCLUDE: ") { + include_dirs.push(PathBuf::from(path)); + } else if let Some(abi) = line.strip_prefix("CXX11_ABI: ") { + cxx11_abi = Some(abi.trim().to_owned()); + } + } + match cxx11_abi { + Some(cxx11_abi) if !include_dirs.is_empty() => Ok(LibtorchConfig { + include_dirs, + cxx11_abi, + }), + _ => Err(format!("unexpected probe output: {stdout}")), + } +} + +fn config_from_libtorch_dir() -> Result { + let root = env_var_rerun("LIBTORCH") + .map(PathBuf::from) + .or_else(|| { + // torch-sys exports the lib dir it linked against (`links = "tch"`). + env::var("DEP_TCH_LIBTORCH_LIB") + .ok() + .map(|lib| PathBuf::from(lib).join("..")) + }) + .ok_or("neither LIBTORCH_USE_PYTORCH nor LIBTORCH is set")?; + let include_dirs = match env_var_rerun("LIBTORCH_INCLUDE") { + Some(include) => vec![ + PathBuf::from(&include).join("include"), + PathBuf::from(&include).join("include/torch/csrc/api/include"), + ], + None => vec![ + root.join("include"), + root.join("include/torch/csrc/api/include"), + ], + }; + Ok(LibtorchConfig { + include_dirs, + cxx11_abi: env_var_rerun("LIBTORCH_CXX11_ABI").unwrap_or_else(|| "1".to_owned()), + }) +} + +fn main() { + println!("cargo:rerun-if-changed=csrc/tensor_from_blob.cpp"); + + let config = if env_var_rerun("LIBTORCH_USE_PYTORCH").is_some() { + config_from_python() + } else { + config_from_libtorch_dir() + } + .unwrap_or_else(|e| panic!("locating libtorch headers for csrc/tensor_from_blob.cpp: {e}")); + + let mut build = cc::Build::new(); + build + .cpp(true) + .pic(true) + .warnings(false) + .includes(&config.include_dirs) + .file("csrc/tensor_from_blob.cpp"); + if env::var("CARGO_CFG_TARGET_OS").as_deref() == Ok("windows") { + build.flag("/std:c++17"); + } else { + build + .flag("-std=c++17") + .flag(format!("-D_GLIBCXX_USE_CXX11_ABI={}", config.cxx11_abi)); + } + build.compile("avtensor_blob"); +} diff --git a/csrc/tensor_from_blob.cpp b/csrc/tensor_from_blob.cpp new file mode 100644 index 0000000..90911d1 --- /dev/null +++ b/csrc/tensor_from_blob.cpp @@ -0,0 +1,46 @@ +// Exposes `at::from_blob` with a deleter to Rust. tch-rs wraps only the +// non-owning overload (`at_tensor_of_blob`); this one lets a buffer allocated +// outside torch's CPU allocator be owned by a tensor and freed by `deleter` +// when the last reference goes away. See src/alloc.rs for why. +#include + +#include +#include +#include +#include +#include + +extern "C" { + +typedef void (*avtensor_blob_deleter)(void *data, void *ctx); + +// Wraps `data` in a contiguous CPU tensor of `sizes` and element type +// `scalar_type` (a c10::ScalarType value). Ownership of `data` moves to the +// tensor's storage: `deleter(data, ctx)` runs when the last tensor sharing +// that storage is released, on whichever thread releases it. +// +// Returns a heap-allocated at::Tensor, the object tch-rs calls `C_tensor` +// and frees with `at_free`. On failure returns nullptr and, if `err` is not +// null, stores a malloc'd message in *err that the caller frees. +void *avtensor_tensor_from_blob(void *data, const int64_t *sizes, size_t ndim, + int32_t scalar_type, + avtensor_blob_deleter deleter, void *ctx, + char **err) { + try { + auto options = at::TensorOptions() + .dtype(static_cast(scalar_type)) + .device(at::kCPU); + at::Tensor tensor = at::from_blob( + data, at::IntArrayRef(sizes, ndim), + [deleter, ctx](void *p) { deleter(p, ctx); }, options); + return new at::Tensor(std::move(tensor)); + } catch (const std::exception &e) { + if (err) *err = strdup(e.what()); + return nullptr; + } catch (...) { + if (err) *err = strdup("unknown exception in at::from_blob"); + return nullptr; + } +} + +} // extern "C" diff --git a/pyproject.toml b/pyproject.toml index 6303d3c..264bc9a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,7 +36,7 @@ auditwheel = "skip" # Include some directives to instruct uv to recompile on changes to Rust files. [tool.uv] -cache-keys = [{file = "pyproject.toml"}, {file = "Cargo.toml"}, {file = "**/*.rs"}] +cache-keys = [{file = "pyproject.toml"}, {file = "Cargo.toml"}, {file = "**/*.rs"}, {file = "csrc/*.cpp"}] [tool.ruff] line-length = 88 diff --git a/src/alloc.rs b/src/alloc.rs new file mode 100644 index 0000000..1313c25 --- /dev/null +++ b/src/alloc.rs @@ -0,0 +1,300 @@ +//! Allocation of the CPU output buffers that decoded frames and samples are +//! written into. +//! +//! By default these buffers come from the system allocator (`posix_memalign`) +//! and are wrapped in a tensor with `at::from_blob` plus a deleter, instead of +//! being allocated with `at::empty`. Shape, dtype, strides and device are +//! identical either way; only the storage's owner differs. +//! +//! Why: torch's CPU allocator is not always plain malloc. The aarch64 Linux +//! wheels embed mimalloc, which keeps a block freed by a thread other than the +//! one that allocated it: the pages are `MADV_FREE`d but stay mapped and +//! resident until the allocating thread exits. Data-loading pipelines decode +//! on worker threads and drop the tensors on other threads, so with `at::empty` +//! every decode left its whole output buffer (hundreds of MiB to GiBs) in +//! resident memory for the lifetime of the worker. glibc releases large blocks +//! with `munmap` on `free` regardless of the calling thread, so the deleter +//! path returns the memory immediately. See `benchmarks/cross_thread_retention.py` +//! for a reproduction. +//! +//! `AVTENSOR_OUTPUT_ALLOCATOR=torch` restores `at::empty` for comparison. + +use std::ffi::{c_char, c_int, c_void, CStr}; +use std::ptr; +use std::sync::OnceLock; + +use anyhow::{anyhow, bail, Context}; +use tch::{Device, Kind, Tensor}; + +/// Which allocator backs CPU output tensors. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OutputAllocator { + /// `posix_memalign` + `at::from_blob` with a `free` deleter (default). + System, + /// `at::empty`, i.e. torch's own CPU allocator. + Torch, +} + +impl OutputAllocator { + fn from_env_value(value: Option<&str>) -> Self { + match value.map(str::trim) { + None | Some("") => OutputAllocator::System, + Some(v) if v.eq_ignore_ascii_case("system") => OutputAllocator::System, + Some(v) if v.eq_ignore_ascii_case("torch") => OutputAllocator::Torch, + Some(v) => { + log::warn!( + "Unknown AVTENSOR_OUTPUT_ALLOCATOR={v:?} (expected 'system' or 'torch'); \ + using 'system'" + ); + OutputAllocator::System + } + } + } +} + +/// The allocator selected by `AVTENSOR_OUTPUT_ALLOCATOR`, read once. +pub fn output_allocator() -> OutputAllocator { + static CHOICE: OnceLock = OnceLock::new(); + *CHOICE.get_or_init(|| { + OutputAllocator::from_env_value(std::env::var("AVTENSOR_OUTPUT_ALLOCATOR").ok().as_deref()) + }) +} + +/// Allocates an uninitialized, contiguous tensor of `shape` and `kind` on +/// `device`. CPU tensors use the allocator selected by [`output_allocator`]; +/// other devices always go through torch. +pub fn empty(shape: &[i64], kind: Kind, device: Device) -> Result { + if device != Device::Cpu || output_allocator() == OutputAllocator::Torch { + return Tensor::f_empty(shape, (kind, device)).map_err(Into::into); + } + empty_cpu_system(shape, kind) +} + +/// Alignment of the system-allocated buffers; matches torch's CPU allocator +/// so vectorized kernels see the same alignment either way. +const ALIGNMENT: usize = 64; + +fn empty_cpu_system(shape: &[i64], kind: Kind) -> Result { + let scalar_type = scalar_type(kind)?; + let numel = shape + .iter() + .try_fold(1usize, |acc, &dim| { + usize::try_from(dim) + .ok() + .and_then(|dim| acc.checked_mul(dim)) + }) + .ok_or_else(|| anyhow!("invalid output shape {shape:?}"))?; + let bytes = numel + .checked_mul(kind.elt_size_in_bytes()) + .ok_or_else(|| anyhow!("output shape {shape:?} overflows usize"))?; + if bytes == 0 { + // Nothing to own; torch handles empty tensors without allocating. + return Tensor::f_empty(shape, (kind, Device::Cpu)).map_err(Into::into); + } + + let mut data: *mut c_void = ptr::null_mut(); + // SAFETY: plain libc call with a valid out-pointer and a power-of-two + // alignment that is a multiple of sizeof(void*). + let rc = unsafe { libc::posix_memalign(&mut data, ALIGNMENT, bytes) }; + if rc != 0 || data.is_null() { + bail!( + "posix_memalign({bytes} bytes) for output shape {shape:?} failed: {}", + std::io::Error::from_raw_os_error(rc) + ); + } + // SAFETY: `data` is a fresh allocation of exactly `bytes` bytes, and + // `free_system_buffer` is the matching release for `posix_memalign`. + unsafe { + tensor_from_blob( + data, + shape, + scalar_type, + free_system_buffer, + ptr::null_mut(), + ) + .with_context(|| format!("wrapping a {bytes}-byte buffer as a {kind:?} tensor")) + } +} + +unsafe extern "C" fn free_system_buffer(data: *mut c_void, _ctx: *mut c_void) { + libc::free(data); +} + +/// `c10::ScalarType` value for `kind` (tch keeps its own mapping private). +fn scalar_type(kind: Kind) -> Result { + Ok(match kind { + Kind::Uint8 => 0, + Kind::Int8 => 1, + Kind::Int16 => 2, + Kind::Int => 3, + Kind::Int64 => 4, + Kind::Half => 5, + Kind::Float => 6, + Kind::Double => 7, + Kind::Bool => 11, + Kind::BFloat16 => 15, + other => bail!("unsupported output dtype {other:?}"), + }) +} + +/// Deleter signature expected by the C++ shim: `(data, ctx)`. +type BlobDeleter = unsafe extern "C" fn(*mut c_void, *mut c_void); + +extern "C" { + fn avtensor_tensor_from_blob( + data: *mut c_void, + sizes: *const i64, + ndim: usize, + scalar_type: i32, + deleter: BlobDeleter, + ctx: *mut c_void, + err: *mut *mut c_char, + ) -> *mut torch_sys::C_tensor; +} + +/// Wraps `data` in a contiguous CPU tensor that owns it: `deleter(data, ctx)` +/// runs when the last tensor sharing the storage is dropped. +/// +/// # Safety +/// +/// `data` must stay valid, and hold `numel(shape) * element size` bytes, until +/// `deleter` is called. `deleter` must be sound to call from any thread. On +/// error the buffer is intentionally leaked: torch may or may not have run the +/// deleter already, and leaking beats a double free on a path that only +/// triggers for invalid shapes. +unsafe fn tensor_from_blob( + data: *mut c_void, + shape: &[i64], + scalar_type: c_int, + deleter: BlobDeleter, + ctx: *mut c_void, +) -> Result { + let mut err: *mut c_char = ptr::null_mut(); + let tensor = avtensor_tensor_from_blob( + data, + shape.as_ptr(), + shape.len(), + scalar_type, + deleter, + ctx, + &mut err, + ); + if tensor.is_null() { + let message = if err.is_null() { + "unknown error".to_owned() + } else { + let message = CStr::from_ptr(err).to_string_lossy().into_owned(); + libc::free(err as *mut c_void); + message + }; + bail!("at::from_blob failed for shape {shape:?}: {message}"); + } + Ok(Tensor::from_ptr(tensor)) +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use super::*; + + #[test] + fn env_value_selects_allocator() { + assert_eq!( + OutputAllocator::from_env_value(None), + OutputAllocator::System + ); + assert_eq!( + OutputAllocator::from_env_value(Some("")), + OutputAllocator::System + ); + assert_eq!( + OutputAllocator::from_env_value(Some("system")), + OutputAllocator::System + ); + assert_eq!( + OutputAllocator::from_env_value(Some("Torch")), + OutputAllocator::Torch + ); + assert_eq!( + OutputAllocator::from_env_value(Some("bogus")), + OutputAllocator::System + ); + } + + #[test] + fn system_buffer_tensor_behaves_like_empty() -> Result<(), anyhow::Error> { + let shape = [3i64, 4, 5, 3]; + let mut tensor = empty_cpu_system(&shape, Kind::Uint8)?; + assert_eq!(tensor.size(), shape); + assert_eq!(tensor.kind(), Kind::Uint8); + assert_eq!(tensor.device(), Device::Cpu); + assert!(tensor.is_contiguous()); + assert_eq!(tensor.data_ptr() as usize % ALIGNMENT, 0); + + // Writable through the usual paths and through views. + let _ = tensor.fill_(1); + let _ = tensor.narrow(0, 1, 1).fill_(5); + assert_eq!(i64::try_from(tensor.sum(Kind::Int64))?, 120 + 60 * 5); + + let mut audio = empty_cpu_system(&[2, 1000], Kind::Float)?; + assert_eq!(audio.kind(), Kind::Float); + let _ = audio.fill_(0.5); + assert!((f64::try_from(audio.sum(Kind::Double))? - 1000.0).abs() < 1e-6); + Ok(()) + } + + #[test] + fn zero_sized_shape_falls_back_to_torch() -> Result<(), anyhow::Error> { + let tensor = empty_cpu_system(&[0, 720, 1280, 3], Kind::Uint8)?; + assert_eq!(tensor.numel(), 0); + assert_eq!(tensor.size(), [0, 720, 1280, 3]); + Ok(()) + } + + #[test] + fn invalid_shapes_are_rejected() { + assert!(empty_cpu_system(&[-1, 3], Kind::Uint8).is_err()); + assert!(empty_cpu_system(&[i64::MAX, 2], Kind::Float).is_err()); + } + + unsafe extern "C" fn counting_deleter(data: *mut c_void, ctx: *mut c_void) { + (*(ctx as *const AtomicUsize)).fetch_add(1, Ordering::SeqCst); + libc::free(data); + } + + #[test] + fn deleter_runs_once_when_the_last_view_drops_on_another_thread() -> Result<(), anyhow::Error> { + static FREED: AtomicUsize = AtomicUsize::new(0); + let shape = [8i64, 16]; + let bytes = 8 * 16 * Kind::Float.elt_size_in_bytes(); + let data = unsafe { libc::malloc(bytes) }; + assert!(!data.is_null()); + let mut tensor = unsafe { + tensor_from_blob( + data, + &shape, + scalar_type(Kind::Float)?, + counting_deleter, + &FREED as *const AtomicUsize as *mut c_void, + )? + }; + let _ = tensor.fill_(2.0); + let view = tensor.narrow(0, 2, 3); + drop(tensor); + assert_eq!( + FREED.load(Ordering::SeqCst), + 0, + "a live view keeps the buffer" + ); + + std::thread::spawn(move || { + assert!((f64::try_from(view.sum(Kind::Double)).unwrap() - 96.0).abs() < 1e-9); + drop(view); + }) + .join() + .unwrap(); + assert_eq!(FREED.load(Ordering::SeqCst), 1); + Ok(()) + } +} diff --git a/src/decoder/mod.rs b/src/decoder/mod.rs index 6bf5ad8..8a7efbf 100644 --- a/src/decoder/mod.rs +++ b/src/decoder/mod.rs @@ -967,7 +967,7 @@ fn init_video_stream_context( .device .map_or(tch::Device::Cpu, |d| tch::Device::Cuda(d as usize)); let dest = - tch::Tensor::f_empty(&plan.shape, (kind, device)).context("allocating output Tensor")?; + crate::alloc::empty(&plan.shape, kind, device).context("allocating output Tensor")?; Ok(StreamContext { stream_type: StreamType::Video, @@ -1202,7 +1202,7 @@ fn init_audio_stream_context( plan.shape, audio_index ); - let dest = tch::Tensor::f_empty(&plan.shape, (tch::Kind::Float, tch::Device::Cpu)) + let dest = crate::alloc::empty(&plan.shape, tch::Kind::Float, tch::Device::Cpu) .context("allocating audio output Tensor")?; Ok(StreamContext { diff --git a/src/lib.rs b/src/lib.rs index 9001d46..4d2b77c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,6 +7,7 @@ use google_cloud_storage::client::Storage; use google_cloud_storage::retry_policy::RetryableErrors; use tokio::runtime::Runtime; +mod alloc; mod conversion; mod decoder; mod ffi; From fed29b8b100d5bffab64d9117d9b9f40b89e9f73 Mon Sep 17 00:00:00 2001 From: EthanRosenthal Date: Sun, 4 Oct 2026 20:52:22 +0000 Subject: [PATCH 2/3] Accept tch/torch-sys 0.24 through 0.26 Each tch release targets a single libtorch version (0.24 for torch 2.11, 0.26 for torch 2.13), and avtensor's source compiles against both. Accept the range in Cargo.toml so a build for torch 2.11 can select tch 0.24 with cargo update --precise instead of editing the manifest. Cargo.lock stays on 0.26, so the default build is unchanged. --- CLAUDE.md | 7 +++++-- Cargo.toml | 8 ++++++-- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index fd3b2cb..b3edd2b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -12,8 +12,11 @@ generated locally with FFmpeg; end-to-end cloud-storage tests live in the consuming project. `tch` comes from crates.io; its version maps 1:1 to a libtorch version -(0.24.0 ↔ torch 2.11), so bump them together, along with the `torch-sys` -pin (same version as `tch`). +(0.24.0 ↔ torch 2.11, 0.26.0 ↔ torch 2.13). `Cargo.toml` accepts a range +so one source tree builds against either torch; `Cargo.lock` holds the +default. To build against another torch, run +`cargo update -p tch -p torch-sys --precise ` (keep `tch` and +`torch-sys` on the same version). `build.rs` compiles `csrc/tensor_from_blob.cpp` against the libtorch headers (located the same way torch-sys locates them); `src/alloc.rs` uses it to give diff --git a/Cargo.toml b/Cargo.toml index 176014f..1ac74a2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,9 +27,13 @@ log = { version = "0.4.27", features = ["kv"] } pyo3 = "0.25.0" pyo3-log = "0.12.4" rsmpeg = { version = "0.17.0", features = ["ffmpeg7_1"] } -tch = { version = "0.26.0", features = ["python-extension"] } +# Each tch release targets one libtorch (0.24 <-> torch 2.11, 0.26 <-> torch 2.13). +# The source builds against both; Cargo.lock picks the default, and a +# build for another torch selects its match with +# `cargo update -p tch -p torch-sys --precise `. +tch = { version = ">=0.24.0, <0.27", features = ["python-extension"] } # Same version as tch; only used for the opaque `C_tensor` handle in src/alloc.rs. -torch-sys = "0.26.0" +torch-sys = ">=0.24.0, <0.27" tempfile = "3.20.0" thiserror = "2.0.15" libloading = "0.8" From d1cdd42ceaebff37005cb66415714cfaf419e4c3 Mon Sep 17 00:00:00 2001 From: EthanRosenthal Date: Sun, 4 Oct 2026 22:04:22 +0000 Subject: [PATCH 3/3] Drop the cross-thread retention benchmark script --- README.md | 3 +- benchmarks/cross_thread_retention.py | 176 --------------------------- src/alloc.rs | 3 +- 3 files changed, 2 insertions(+), 180 deletions(-) delete mode 100644 benchmarks/cross_thread_retention.py diff --git a/README.md b/README.md index 976bdf8..5dba5f0 100644 --- a/README.md +++ b/README.md @@ -423,8 +423,7 @@ provider-specific support — the URL is handed to FFmpeg's http protocol. exits). avtensor therefore allocates its output buffers with the system allocator and hands them to torch via `from_blob` with a deleter, so a free is an immediate `munmap` on any thread. `AVTENSOR_OUTPUT_ALLOCATOR=torch` - switches back to `torch.empty`-style allocation; - `benchmarks/cross_thread_retention.py` measures the difference. + switches back to `torch.empty`-style allocation. ## Contributing diff --git a/benchmarks/cross_thread_retention.py b/benchmarks/cross_thread_retention.py deleted file mode 100644 index 8decd9e..0000000 --- a/benchmarks/cross_thread_retention.py +++ /dev/null @@ -1,176 +0,0 @@ -"""Measure memory retained after decoded tensors are freed on another thread. - -Decodes one clip repeatedly, first allocating and dropping on the main thread, -then decoding on worker threads and dropping on the main thread (the shape of a -typical data-loading pipeline), and reports the process's resident memory after -each phase from ``/proc/self/smaps_rollup``. ``LazyFree`` counts pages that were -``MADV_FREE``'d but are still resident, which is how a thread-caching allocator -keeps a cross-thread free alive. - -Run it twice to compare the output allocators:: - - python benchmarks/cross_thread_retention.py --clip clip.mp4 - AVTENSOR_OUTPUT_ALLOCATOR=torch python benchmarks/cross_thread_retention.py --clip clip.mp4 - -Without ``--clip`` a synthetic 1280x720 H.264 clip with audio is generated with -``ffmpeg`` (``$FFMPEG`` or the one on ``PATH``). Linux only (reads ``/proc``). -""" - -from __future__ import annotations - -import argparse -import gc -import os -import shutil -import subprocess -import sys -import tempfile -import time -from concurrent.futures import ThreadPoolExecutor -from pathlib import Path - -import avtensor -import torch - -GIB = 2**30 - - -def smaps_rollup() -> dict[str, float]: - """Rss / Anonymous / LazyFree of this process in GiB.""" - wanted = {"Rss", "Anonymous", "LazyFree"} - out: dict[str, float] = dict.fromkeys(wanted, 0.0) - with open("/proc/self/smaps_rollup") as f: - for line in f: - key, _, rest = line.partition(":") - if key in wanted: - out[key] = int(rest.split()[0]) * 1024 / GIB - return out - - -def generate_clip(path: Path, seconds: float, size: str, fps: int) -> None: - ffmpeg = os.environ.get("FFMPEG") or shutil.which("ffmpeg") - if ffmpeg is None: - sys.exit("ffmpeg not found (set FFMPEG or PATH), or pass --clip instead") - cmd = [ - ffmpeg, - "-hide_banner", - "-loglevel", - "error", - "-y", - "-f", - "lavfi", - "-i", - f"testsrc2=size={size}:rate={fps}:duration={seconds}", - "-f", - "lavfi", - "-i", - f"sine=frequency=440:sample_rate=48000:duration={seconds}", - "-c:v", - "libx264", - "-preset", - "ultrafast", - "-pix_fmt", - "yuv420p", - "-c:a", - "aac", - str(path), - ] - subprocess.run(cmd, check=True) - - -def decode(clip: str, with_audio: bool) -> list[torch.Tensor]: - """Decode the clip and return its output tensors (so the caller frees them).""" - request = avtensor.MediaDecodeRequest( - clip, - video_stream=avtensor.VideoStreamRequest(), - audio_streams=[avtensor.AudioStreamRequest()] if with_audio else None, - ) - return [s["data"] for s in avtensor.decode_asset(request)] - - -def nbytes(tensors: list[torch.Tensor]) -> int: - return sum(t.numel() * t.element_size() for t in tensors) - - -def report( - tag: str, base: dict[str, float], times: list[float], bytes_out: int -) -> None: - m = smaps_rollup() - mean_ms = 1000 * sum(times) / len(times) if times else float("nan") - print( - f"{tag:<34} rss={m['Rss']:6.2f} anon={m['Anonymous']:6.2f} " - f"lazyfree={m['LazyFree']:6.2f} GiB " - f"(+{m['Rss'] - base['Rss']:.2f} GiB over start) " - f"decode {mean_ms:7.1f} ms output {bytes_out / GIB:.2f} GiB/decode", - flush=True, - ) - - -def main() -> None: - parser = argparse.ArgumentParser( - description="Measure memory retained after cross-thread frees of decoded tensors." - ) - parser.add_argument("--clip", type=Path, help="media file to decode") - parser.add_argument("--decodes", type=int, default=12, help="decodes per phase") - parser.add_argument("--threads", type=int, default=4, help="worker threads") - parser.add_argument( - "--seconds", type=float, default=20.0, help="synthetic clip length" - ) - parser.add_argument("--size", default="1280x720", help="synthetic clip size") - parser.add_argument("--fps", type=int, default=24, help="synthetic clip frame rate") - parser.add_argument( - "--settle", type=float, default=0.3, help="seconds to wait after each free" - ) - args = parser.parse_args() - - allocator = os.environ.get("AVTENSOR_OUTPUT_ALLOCATOR", "system (default)") - print(f"AVTENSOR_OUTPUT_ALLOCATOR={allocator}", flush=True) - - tmpdir: tempfile.TemporaryDirectory[str] | None = None - clip = args.clip - if clip is None: - tmpdir = tempfile.TemporaryDirectory() - clip = Path(tmpdir.name) / "clip.mp4" - generate_clip(clip, args.seconds, args.size, args.fps) - clip = str(clip) - - with_audio = bool(avtensor.probe_asset(clip)["audio_streams"]) - - # Warm up (codec init, allocator arenas) so the baseline is steady. - bytes_out = nbytes(decode(clip, with_audio)) - gc.collect() - base = smaps_rollup() - report("start", base, [], bytes_out) - - times: list[float] = [] - for _ in range(args.decodes): - t0 = time.monotonic() - tensors = decode(clip, with_audio) - times.append(time.monotonic() - t0) - del tensors - time.sleep(args.settle) - report(f"{args.decodes}x decode+free on main thread", base, times, bytes_out) - - pool = ThreadPoolExecutor(max_workers=args.threads) - times = [] - for _ in range(args.decodes): - t0 = time.monotonic() - tensors = pool.submit(decode, clip, with_audio).result() - times.append(time.monotonic() - t0) - del tensors # freed on the main thread, not the decoding thread - time.sleep(args.settle) - report(f"{args.decodes}x decode in pool, free on main", base, times, bytes_out) - - time.sleep(5.0) - report("after 5 s idle", base, times, bytes_out) - - pool.shutdown(wait=True) - time.sleep(0.5) - report("after worker threads exited", base, times, bytes_out) - - if tmpdir is not None: - tmpdir.cleanup() - - -if __name__ == "__main__": - main() diff --git a/src/alloc.rs b/src/alloc.rs index 1313c25..3b7bf2d 100644 --- a/src/alloc.rs +++ b/src/alloc.rs @@ -14,8 +14,7 @@ //! every decode left its whole output buffer (hundreds of MiB to GiBs) in //! resident memory for the lifetime of the worker. glibc releases large blocks //! with `munmap` on `free` regardless of the calling thread, so the deleter -//! path returns the memory immediately. See `benchmarks/cross_thread_retention.py` -//! for a reproduction. +//! path returns the memory immediately. //! //! `AVTENSOR_OUTPUT_ALLOCATOR=torch` restores `at::empty` for comparison.