Skip to content

feat(python): bound dataset runtime resources and add trusted attestations - #950

Draft
RkGrit wants to merge 3 commits into
feature/trusted-index-modefrom
gewu/timer4-batched-window-read-clean
Draft

RkGrit wants to merge 3 commits into
feature/trusted-index-modefrom
gewu/timer4-batched-window-read-clean

Conversation

@RkGrit

@RkGrit RkGrit commented Sep 9, 2026

Copy link
Copy Markdown
Contributor

Summary

This draft extends feature/trusted-index-mode with the resource-lifetime and identity contracts needed by large, multi-process training workloads:

  1. release the original Dataset Index file object immediately after a successful read-only mmap on Linux;
  2. make DatasetRuntime the explicit owner of index mappings, Reader sessions, prepared-series handles, the query executor, and their leases;
  3. bound Reader and prepared-series retention, including failure, close, and fork() paths;
  4. add immutable index/data-file attestations so one trusted parent can perform the expensive identity pass once and workers can reopen the exact immutable dataset without rescanning directories or rehashing the index;
  5. bind prepared-series cache entries to versioned file-generation identities rather than Python object identity.

The changes are confined to the Python Dataset API. They do not add TimerLathe-specific imports, change the C++ reader ABI, implement native fused multi-window I/O, or claim end-to-end Arrow parity.

Downstream integration: https://github.com/thulab/TimerLathe/pull/6

Base and scope

  • Base: feature/trusted-index-mode@4159c635de8f29c973f56d06a11771c8c0121ff9
  • Head: gewu/timer4-batched-window-read-clean@94649f3e4a1d0f1245cf0616ac24c3853db2769a
  • Commits:
    • e9f28e09 — release the Dataset Index file after mmap;
    • 3d21b820 — bound runtime resource ownership and harden fork()/failure cleanup;
    • 94649f3e — add trusted index/data attestations and generation identities.
  • Diff: 7 Python files, 2,418 insertions, 115 deletions, including extensive lifecycle and failure-injection tests.

Motivation

The motivating downstream workload opens a fixed TimeBench asset containing 218 logical datasets from 8 distributed ranks with multiple DataLoader workers per rank. The dataset is immutable for the duration of training and already has trusted .tsidx files. At this scale, a resource cost that looks small for one TsFileDataFrame becomes multiplicative:

rank × DataLoader worker × live TsFileDataFrame × retained Reader/index resource

The existing trusted-index mode removes expensive validation work, but it did not yet define a complete handoff from a trusted parent process to workers, nor an explicit lifecycle for every native/index resource. The concrete failure modes were:

  • an index mapping could retain both the original Python file object and the descriptor duplicated by CPython's mmap implementation;
  • cached Readers and prepared-series handles needed hard bounds and precise active-use semantics;
  • close/error paths could stop after the first exception and leak later-owned resources;
  • DataLoader fork() could inherit locks, thread-pool state, Readers, prepared handles, and leases that are unsafe to reuse in the child;
  • a worker needed a cheap way to prove that the index and data files still represent the exact asset approved by the parent, without repeating full hashes or directory scans.

This PR treats those as one ownership problem: expensive identity work happens once at the trust boundary; each process owns only its own mutable/native runtime; every retained object has an explicit cap, lease, and deterministic teardown path.

Design overview

flowchart LR
    P[Trusted parent] --> IH[Hash finalized .tsidx once]
    P --> DF[Open and fstat data files]
    IH --> IA[IndexAttestation]
    DF --> DA[Ordered DataFileAttestations]
    IA --> W[Worker open]
    DA --> W
    W --> V[Canonical path + stat validation]
    V --> M[Read-only mapped index]
    M --> R[Bounded ReaderSessionPool]
    M --> C[Bounded PreparedSeriesCache]
    R --> Q[Queries]
    C --> Q
Loading

The parent hashes index bytes only once. Workers validate canonical path and stable file metadata, map the attested index, and create fresh process-local Reader/cache/executor state. Mismatch fails before value I/O.

Detailed changes

1. Drop the original .tsidx file descriptor after mmap

MappedDatasetIndex now initializes every owned resource to None, creates the read-only mapping, and closes the original file object immediately on Linux. The mapping and its memoryview remain live, so index access is unchanged.

The important distinction is:

  • the original open(..., "rb") descriptor is released;
  • CPython's mmap object still owns one descriptor on the tested runtime;
  • the VMA remains until MappedDatasetIndex.close();
  • closing the index releases the view, mapping, and any remaining file object in order.

_close_owned_resources() uses nested finally blocks, is idempotent, and continues cleaning later resources even if an earlier release raises. Constructor failures before/after mapping use the same path.

This is intentionally narrower than a “zero-FD mmap” or an end-to-end zero-copy claim.

2. Explicit runtime leases and quiescent teardown

DatasetRuntime now owns:

  • MappedDatasetIndex;
  • ReaderSessionPool;
  • PreparedSeriesCache;
  • the optional query ThreadPoolExecutor;
  • the mapped catalog exposed to TsFileDataFrame/Timeseries views.

RuntimeLease protects an object/view lifetime. _QueryLease protects an active query. When the last object lease closes, the runtime stops accepting work, waits for active query leases to drain, and tears resources down exactly once.

This prevents a subset view or result path from closing native state still in use by another view. It also gives close/failure code a single ownership graph instead of several loosely related dictionaries.

3. Hard-bounded Reader sessions

ReaderSessionPool is a per-runtime LRU with an explicit max_open_files cap:

  • an active Reader is never evicted;
  • an idle LRU Reader is closed before opening its replacement;
  • if all slots are active, acquisition waits instead of exceeding the cap;
  • close() waits for active uses to finish and then closes all Readers;
  • cleanup is idempotent and continues past individual close failures.

The existing TSFILE_DATAFRAME_MAX_OPEN_FILES setting controls this cap (default 16 in DatasetRuntime). This is a per-runtime/per-DataFrame bound, not yet a global per-process FD budget. Applications retaining many DataFrames must still size the product of frame count and per-frame Reader cap.

4. Bounded, generation-aware PreparedSeriesCache

The prepared-series cache is now a runtime-wide, single-flight LRU:

  • TSFILE_DATAFRAME_MAX_PREPARED_SERIES sets a non-negative retention cap;
  • 0 is a supported “no idle retention” mode;
  • concurrent misses for the same key prepare once;
  • active entries are not evicted;
  • aligned-series time-owner dependencies are tracked explicitly, so an owner cannot be closed before dependents;
  • eviction and shutdown close every prepared handle exactly once;
  • BaseException, decode failure, partial construction, and close failure paths release leases and continue cleanup.

With attestations, cache keys use file_generation_token(index_identity, data_manifest_identity, file) plus locator ID. Without attestations, the previous mapped-index/file metadata identity remains the compatibility fallback.

5. Fork-safe process-local runtime replacement

Locks, native Readers, prepared handles, and executor threads are not reused across fork().

  • leases record their creator PID and reject inherited query use;
  • inherited resources can be discarded in the child without taking parent-created locks;
  • TsFileDataFrame._ensure_process_local_runtime() creates a fresh child runtime on first use;
  • subset views rebind to the root's child-local runtime when possible;
  • parent resources remain usable and independently closeable;
  • repeated child cleanup is safe.

This targets PyTorch/DataLoader-style worker creation while keeping normal single-process behavior unchanged.

6. Public trusted-attestation API

The new generic API exported from tsfile.dataset is:

API Purpose
IndexAttestation Immutable schema/version/path/content/stat identity for one finalized .tsidx
DataFileAttestation Ordered file ID, canonical path, stat identity, and optional existing index fingerprint
attest_index() Hash the opened index bytes once and reject mutation during hashing
validate_index_attestation() Worker-side canonical-path/stat check without rescanning or rehashing bytes
attest_data_file() Capture identity from the opened data-file object via fstat
validate_data_file_attestation() Validate the worker-visible path/stat identity
data_manifest_identity() Produce a versioned, length-prefixed, order-sensitive manifest digest
file_generation_token() Bind one file to the index digest and complete ordered data manifest
AttestationMismatchError Fail-closed mismatch reported before value reads

TsFileDataFrame and DatasetRuntime accept index_attestation and data_file_attestations only as a complete pair and only with trust_index=True. In the attested path:

  • paths come from the ordered attestations, avoiding directory expansion/scans;
  • file IDs must cover 0..N-1 in exact order;
  • indexed path, size, and existing fingerprint must agree;
  • the index is validated before mapping and the mapped object's identity is checked again;
  • each data file is validated before its native Reader opens;
  • any mismatch fails closed.

The encoding uses explicit domain separators, lengths, little-endian uint64 fields, and a schema version so it does not depend on pickle, locale, mapping order, or Python object identity.

Trust and security boundary

Attested trusted mode is for a prebuilt, immutable dataset managed by the application. It does not turn untrusted input into trusted input and does not continuously monitor files after open.

The intended contract is:

  1. the owner verifies/finalizes the asset and performs the expensive index hash once;
  2. the asset remains immutable for the job;
  3. workers cheaply validate the exact canonical path/stat/fingerprint tuple;
  4. mismatch aborts before querying.

Safe mode remains the default. Callers that do not pass attestations keep existing path expansion, index matching, validation, and generation checks.

Observable effects

Resource behavior proven by tests

  • A mapped index retains one target-inode FD and one live VMA on the tested Linux/CPython runtime; the original Python file object is gone.
  • Closing returns both target-inode FD and VMA counts to baseline.
  • Sixteen simultaneous mapped indexes retain one target-inode FD per mapping rather than retaining the original file object as well.
  • Reader sessions never exceed their configured per-runtime cap and active Readers are never evicted.
  • Prepared-series retention plateaus at its configured cap, including 0, while active/dependent entries remain valid.
  • Parent and child runtimes close independently after fork().
  • Constructor, query, decode, and cleanup fault injection leaves no later-owned resource unclosed.

Performance interpretation

This branch removes repeated identity work from workers and bounds resource growth, but no standalone wall-time benchmark has yet isolated these three commits. It would be misleading to assign the downstream training speedup to this PR alone.

In the downstream Timer4 experiment tree, the larger composed optimization stack observed a short-run data-handoff P50 change from 886.022 ms to 2.348 ms, P95 from 957.214 ms to 4.701 ms, and throughput from 2,835.81 to 4,160.13 samples/s. Those numbers:

  • include TimerLathe-side request blocking, shared compact plans, scalar locator reads, rank-side prefetch, and collation changes;
  • were not reproduced on the clean downstream branch plus a released wheel;
  • did not establish long-run memory stability: a later 2,100-step run had 90.533% of waits below 5 ms, roughly 7% throughput decay, and slowly growing PSS.

They are motivation for further profiling, not a benchmark claim for this PR.

Compatibility and non-goals

  • Safe mode remains the default; attestation arguments are optional.
  • Existing trust_index=True callers without attestations continue to work.
  • No TimerLathe-specific type or import is added.
  • No C/C++ ABI or TsFile on-disk format is changed.
  • No native query_windows_many()/merged chunk-page read is implemented.
  • map_query_groups() still schedules independent logical queries; it is not fused physical I/O.
  • No claim is made that mmap is end-to-end zero-copy.
  • No released wheel artifact is produced by this PR.

Validation

Fresh verification on Linux/Python 3.11:

  • complete python/tests suite: 293 passed;
  • directly related index/runtime/identity suites: 175 passed;
  • all 7 changed Python files pass Ruff formatting;
  • all 7 changed Python files parse successfully;
  • git diff --check passes;
  • 7 Ruff lint findings are unchanged from the base branch; this PR adds zero new lint findings.

Tests cover FD/VMA attribution via /proc, idempotent close, mmap and post-mmap failures, active Reader eviction, cache plateau/single-flight/dependency ordering, BaseException cleanup, process-local reopen after fork(), attestation mutation/mismatch, canonical encoding, Python 3.9-compatible imports, and public API exposure.

Suggested follow-up work

This PR intentionally establishes a safe baseline that TsFile maintainers can profile and extend. The highest-value next steps are:

  1. Publish a reproducible wheel artifact. Build from the reviewed source, record an honest package version, source commit, URL, and SHA-256, and test the installed artifact rather than a source overlay. The downstream TimerLathe recipe currently pins an older 2.5.0.dev0 wheel that does not contain these APIs.
  2. Measure the real process-wide budget. The Reader cap is per runtime. Profile active DataFrames, open Readers, index mmap FDs/VMAs, prepared entries, PSS, and page faults over a long multi-worker run. If needed, add cooperative process-level admission/accounting without forcing application-specific singleton semantics into the API.
  3. Separate Reader object reuse from descriptor retention. Explore whether a reusable parsed Reader can release or lazily reacquire its underlying FD, and how this interacts with the optional mmap backend.
  4. Add native batch-window reading. Group requests by file/device/series/chunk/page, merge compatible physical ranges, decode shared pages once, and restore logical request order. The current downstream “batch” path mainly batches dispatch and reuses Readers.
  5. Advance mmap toward true borrowed-buffer semantics. PR [Feature][C++] Support mmap as an optional local file read backend #920 removes repeated positioned reads but its ReadFile::read() path still copies mapped bytes into a caller buffer. End-to-end zero-copy requires explicit borrowed-span ownership from file mapping through page parsing/decompression/decoding, with safe mapping lifetime and mutation rules.
  6. Expose low-overhead counters. Useful contracts include Reader open/hit/evict, prepared hit/miss/evict, physical query count, decoded bytes/rows, mmap VMA/FD ownership, and query/decode latency. Keep collection optional so training hot paths stay unaffected.
  7. Revisit index scale. For hundreds of datasets, measure VMA/page-table and mapped-index working-set costs independently of Reader/data-page costs before adding broader caches.

Review guide

The most important review questions are:

  • Is every native/index resource owned by exactly one runtime path and closed once on success, failure, and fork()?
  • Are the per-runtime Reader/cache bounds and active-use rules correct under concurrency?
  • Is the attestation encoding stable and sufficiently generic for non-TimerLathe callers?
  • Is the immutable trusted-dataset assumption explicit enough, without implying protection against hostile concurrent file replacement?
  • Should any public names or environment variables be adjusted before a wheel is released?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant