feat(streaming): add EAGLE-3 producer and loader over the local store#3085
Open
kashif wants to merge 4 commits into
Open
feat(streaming): add EAGLE-3 producer and loader over the local store#3085kashif wants to merge 4 commits into
kashif wants to merge 4 commits into
Conversation
This is PR 1 of issue NVIDIA-NeMo#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>
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>
Four review fixes on the SampleRefQueue, all surfaced by @khazic on NVIDIA-NeMo#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>
PR 2 of NVIDIA-NeMo#3062. Closes the minimal end-to-end loop on top of PR 1's data-plane contracts: a colocated target forward through a producer + queue + in-process store, materialized by a loader that yields Eagle3TargetBatch to the trainer. No recipe / dispatch changes deliberately (per the RFC's 'library layer first' Q1 answer). Adds: - streaming/eagle3.py: the EAGLE-3 schema (core features + exactly-one supervision encoding) and the helpers that map an Eagle3TargetBatch into the producer's tensors + per-feature FeatureSpec dict. The consumer-side validate_eagle3_ref is the gate the loader runs. - streaming/producer.py: FeatureProducer wraps the existing colocated HFEagle3TargetModel, runs its generate_batch, packs the tensors under the schema, and returns a tensor-free SampleRef. Algorithm is inferred from the backend's runtime type ('*Eagle3TargetModel') with an explicit override. Rejects a backend that emits the precomputed draft-vocab encoding (the colocated producer only carries full logits; the draft-vocab producer lands alongside an SGLang / remote backend). - streaming/loader.py: FeatureDataLoader is a Python iterator over Eagle3TargetBatch. Releases the previous batch's lease + store handle on the NEXT pull (so the trainer can hold a batch across one forward without it being freed mid-forward). consume_now() releases eagerly after backward. close()/context-manager releases on iteration end. Refs whose algorithm does not match the loader's are failed back to the queue rather than silently dropped. Tests (34 new, all passing on top of PR 1's 60): - test_eagle3.py: schema validator + tensor/feature-spec packing. - test_producer.py: round-trip, algorithm inference, draft-vocab encoding rejection, pass-throughs. - test_loader.py: release-previous-on-next, consume_now, close, algorithm-mismatch redelivery. - test_eagle3_parity.py: bit-identical Eagle3TargetBatch across HFEagle3TargetModel.generate_batch (colocated) and FeatureProducer + FeatureDataLoader (streaming) over the same input. This is the numerical-equivalence guarantee PR 2 lands on. ruff format + ruff check clean. Wider sweep: 859 passed, 12 skipped on tests/unit_tests/speculative/. Zero regressions. Signed-off-by: Kashif Rasul <kashif.rasul@gmail.com>
kashif
force-pushed
the
kashif/feat/streaming-eagle3-producer-pr2
branch
from
279fea9 to
fecb6b1
Compare
3 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR 2 of #3062. Closes the minimal end-to-end loop on top of #3084: a colocated target forward through a producer + queue + LocalFeatureStore, materialized by a loader that yields
Eagle3TargetBatchto the trainer. No recipe or dispatch changes (per the RFC's "library layer first" answer to Q1).What does this PR do ?
PR 2 of #3062. Adds the EAGLE-3 producer/loader pair on top of PR 1's data-plane contracts and verifies numerical parity against the existing colocated
HFEagle3TargetModel.generate_batchpath.Changelog
nemo_automodel/components/speculative/streaming/modules:eagle3.py: the EAGLE-3 schema (core features + exactly-one supervision encoding) plus the helpers that map anEagle3TargetBatchinto the producer's tensors + per-featureFeatureSpecdict.validate_eagle3_refis the gate the loader runs.producer.py:FeatureProducerwraps the existing colocatedHFEagle3TargetModel, runs itsgenerate_batch, packs the tensors under the schema, and returns a tensor-freeSampleRef. Algorithm is inferred from the backend's runtime type (*Eagle3TargetModel) with an explicit override. Rejects a backend that emits the precomputed draft-vocab encoding (the colocated producer only carries full logits).loader.py:FeatureDataLoaderis a Python iterator overEagle3TargetBatch. Releases the previous batch's lease + store handle on the NEXT pull (so the trainer can hold a batch across one forward without it being freed mid-forward).consume_now()releases eagerly afterbackward().close()/ context manager releases on iteration end. Refs whose algorithm does not match the loader's are failed back to the queue rather than silently dropped.tests/unit_tests/speculative/streaming/:test_eagle3.py: schema validator + tensor / feature-spec packing.test_producer.py: round-trip, algorithm inference, draft-vocab encoding rejection, pass-throughs.test_loader.py: release-previous-on-next, consume_now, close, algorithm-mismatch redelivery.test_eagle3_parity.py: bit-identicalEagle3TargetBatchacrossHFEagle3TargetModel.generate_batch(colocated) andFeatureProducer + FeatureDataLoader(streaming) over the same input. The numerical-equivalence guarantee PR 2 lands on.ruff formatclean,ruff checkclean.859 passed, 12 skippedacrosstests/unit_tests/speculative/. Zero regressions.Before your PR is "Ready for review"
Additional Information
target_model_backend: streaming, packed-sequence producer metadata threading, async prefetch, draft-vocab encoding producer.