SIGN IN SIGN UP

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