feat(streaming): add EAGLE-3 producer and loader over the local store (#3085)
* feat(streaming): add data-plane contracts and local feature store This is PR 1 of issue #3062 -- the first 1/4 of the streaming producer / consumer pipeline for EAGLE-3 / DFlash / DSpark draft training. No trainer wiring; pure library + contract tests; behavior-neutral. Adds: - `SampleRef` + `FeatureSpec`: frozen, tensor-free references for the control plane. Validation rejects tensors at construction time and the per-algorithm required-features set keeps a misconfigured producer from slipping a stale key set past the consumer. - `assert_no_tensors`: structural (not nominal) guard against tensors, numpy arrays, and duck-typed tensor-likes anywhere in the dataclass / dict / list tree. - `FeatureStore` ABC + `StoreHandle` + `StoreHealth`: 6-method data plane contract with consume-once refcounting, ints-only health snapshot, and a URI partition so cross-process / cross-store refs are rejected at materialization time. - `LocalFeatureStore`: in-process implementation with sample-count and byte caps (bytes as the hard backstop per RFC Q2) and high / low-watermark hysteresis. - `SampleRefQueue` + `Lease` + `VisibilityTimeout`: lease / ack / fail queue with visibility-timeout reclaim and watermark-based backpressure driven by `StoreHealth` ints only. Tests (60, all passing): - refs: invariants, frozen, algorithm-requirements, assert_no_tensors structural walk, tensor-free construction guard. - local_store: round-trip, detached copies, multi-handle refcount, residency caps (sample + bytes), watermark transitions, lifecycle. - queue: FIFO acquire / ack, fail redelivery, visibility-timeout reclaim, pause / resume callback hysteresis, duplicate-id rejection. Signed-off-by: Kashif Rasul <kashif.rasul@gmail.com> * docs(streaming): refresh copyright year + tensor-contract docstrings Updates the linting skill's guidance (current year for new files) and removes an outdated comment about a 1.0 second VisibilityTimeout minimum (only > 0 is enforced). Strengthens put/get docstrings with explicit ownership / aliasing guarantees so future tensor-lifecycle reviewers have one source of truth. Signed-off-by: Kashif Rasul <kashif.rasul@gmail.com> * fix(streaming): honor review feedback on PR 1 queue backpressure Four review fixes on the SampleRefQueue, all surfaced by @khazic on #3084. Behavior matches the standard sliding-window / message-broker patterns: HWM/LWM hysteresis with state preserved in the band (matches TCP-style flow control + AWS SQS visibility timeout semantics), prompt counter cleanup on terminal ack (matches SQS's 'prompt deletion' recommendation), and a public close-signal property to mirror the queue.Queue empty/closed separation in the Python stdlib. 1. Backpressure hysteresis (queue.py: SampleRefQueue.put_blocks_until_below) Resume previously fired on 'not high_watermark_hit', which is 'resident below high' -- not the documented 'resident at or below low'. The high-low band was therefore dead: a producer sitting at resident=high-1 flapped in and out on every consumer release. Resume now gates on 'low_watermark_hit'; pause still gates on 'high_watermark_hit'; in the band the producer's existing paused/unpaused state is preserved. 2. Wire high_watermark_bytes / low_watermark_bytes ctor args These were stored on self._high_bytes / self._low_bytes but never read; backpressure came entirely from the store's high_watermark_hit. The queue now reads its own ctor values first (resident >= high_bytes to pause, resident <= low_bytes to resume) and falls back to the store's health() booleans when the ctor args are None. Misconfigured pairs (low >= high) raise ValueError at construction. 3. Pop _sample_counters on terminal ack The dict was setdefault-ed on every put but never removed. Long streaming runs grew one entry per unique sample_id forever. Now popped in ack() once the lease is removed from outstanding; fail and reclaim_expired re-set the counter when they re-enqueue, so the redelivery bookkeeping stays consistent. 4. Expose is_closed property acquire() returned None for both 'transient empty poll, retry' and 'shutdown, drained' -- consumers couldn't tell them apart. The property returns the queue's closed flag, so a consumer can do 'if lease is None and q.is_closed: break; if lease is None: continue'. Mirrors the queue.Queue empty/closed separation. Tests: - test_queue_pause_and_resume_callbacks_fire_on_watermark_transitions: drainer now releases both warm and p1 (not just p1) so resident drops strictly below low_watermark_bytes; the old assertion 'drops to ~32 KiB which equals low' silently relied on the broken 'resume at not high' behavior. - 6 new tests in test_queue.py covering: * is_closed property disambiguating empty-poll from shutdown * ack() drops _sample_counters entries (leak fix) * fail() preserves the counter until the terminal ack * explicit ctor args override the store's defaults * ctor arg validation (low < high, both positive) * hysteresis preserves state in the band ruff format + ruff check clean. 18 / 18 test_queue tests pass; 66 / 66 streaming tests pass; 831 passed / 12 skipped wider sweep. Signed-off-by: Kashif Rasul <kashif.rasul@gmail.com> * fix(streaming): lease + handle identity, ref deep-freeze, put duplicate check Round 2 review fixes from @khazic on #3084. PR 2 and PR 3 rebase onto this commit to pick the lease/handle identity work up automatically. 1. Lease identity (queue.py: Lease, ack, fail, reclaim_expired) Lease now carries a per-acquire unique lease_id (module-level counter). _outstanding is keyed by lease_id and ack/fail verify the identity before any state mutation. A late ACK for a reclaimed-then-re-leased ref used to pop the new consumer's active lease for the same sample_id (the previous dict was keyed by sample_id). Identity check makes that stale ack a no-op. 2. put_blocks_until_below duplicate check (queue.py) The non-blocking put() checked for duplicate sample_ids atomically; the blocking put_blocks_until_below did not. A producer that retries a ref used to be able to double-enqueue and the second acquire would overwrite the first lease. Added the same atomic check inside the lock in both resume and put branches. 3. StoreHandle identity (store.py, stores/local.py) StoreHandle now carries a per-get unique handle_id (module-level counter). LocalFeatureStore.release is now keyed by handle_id, not by sample_id: re-releasing the same handle is a true no-op (per-handle dispatch), and one handle's re-release no longer decrements a sibling handle's outstanding count. Storage unlinks only when the last sibling handle goes away. 4. SampleRef deep-freeze (refs.py) frozen=True on the dataclass does not deep-freeze feature_keys or feature_specs -- callers could mutate them after construction, including inserting a torch.Tensor, which would silently break the advertised immutable tensor-free control-plane contract. __post_init__ now wraps both fields in MappingProxyType(dict(...)) so subsequent item assignment / deletion raises TypeError. Tests (12 new): - test_lease_id_is_unique_per_acquire - test_stale_ack_does_not_pop_a_newer_active_lease - test_stale_fail_does_not_pop_a_newer_active_lease - test_ack_for_unknown_lease_id_is_noop - test_put_blocks_until_below_rejects_duplicate_sample_id - test_put_blocks_until_below_rejects_duplicate_while_outstanding - test_local_store_handles_for_same_sample_have_distinct_identities - test_local_store_releasing_one_handle_twice_does_not_decrement_sibling - test_local_store_release_unknown_handle_id_is_noop - test_sample_ref_feature_keys_is_read_only_view - test_sample_ref_feature_specs_is_read_only_view - test_sample_ref_rejects_tensor_in_feature_keys_after_construction ruff format + ruff check clean. 78 / 78 streaming tests pass. Signed-off-by: Kashif Rasul <kashif.rasul@gmail.com> * fix(streaming): drop stale handle_refs put-seed leak and duplicate lease field Signed-off-by: Kashif Rasul <kashif.rasul@gmail.com> * fix(streaming): stable store_uri via uuid, clarify acquire drain Signed-off-by: Kashif Rasul <kashif.rasul@gmail.com> * feat(streaming): eagle3 producer and loader over the local store Signed-off-by: Kashif Rasul <kashif.rasul@gmail.com> * fix(streaming): drop poisoned refs, evict sample on spec mismatch Signed-off-by: Kashif Rasul <kashif.rasul@gmail.com> --------- Signed-off-by: Kashif Rasul <kashif.rasul@gmail.com> Co-authored-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com> Co-authored-by: Huiying <willwin.lee@gmail.com>
K
Kashif Rasul committed
e105bf779bcdf26e41f613bab1bfe92a8eaf13fa
Parent: 4e5c038
Committed by GitHub <noreply@github.com>
on 8/9/2026, 6:23:13 PM