diff --git a/CLAUDE.md b/CLAUDE.md index 50c280b..b3edd2b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -12,7 +12,16 @@ 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, 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 +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..1ac74a2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,7 +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.24.0, <0.27" tempfile = "3.20.0" thiserror = "2.0.15" libloading = "0.8" @@ -38,6 +44,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..5dba5f0 100644 --- a/README.md +++ b/README.md @@ -416,6 +416,14 @@ 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. ## Contributing @@ -423,9 +431,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/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..3b7bf2d --- /dev/null +++ b/src/alloc.rs @@ -0,0 +1,299 @@ +//! 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. +//! +//! `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;