Ref-based object store with dict trace and three-layer architecture

[cweniger]

Summary

• Replace RVBatch with plain dict returns ({'value': ndarray, 'log_prob': ndarray})
• Flat store keys: theta.value, theta.log_prob instead of bare theta / theta.logprob
• Ref-only outer layer: DeployedGraph and store never touch arrays, only ObjectRefs
• NodeWrapper handles ref↔array boundary with batched ray.get / ray.put
• Remove RV, RVBatch, as_rv, as_rvbatch from utils.py
• Remove unnecessary pause/resume during resimulation (async actors interleave naturally)

Performance optimizations

Broadcast detection and deferred expansion

Profiling on 05_linear_regression (n=1024 resample batch) identified bottlenecks in sample_proposal that block the training event loop:

Bottleneck	Before	After	Improvement
_resolve_refs (broadcast observation resolved N times)	280ms	~1ms	~280x — detect identical refs, resolve once, return (1,...) shape
_batch_to_refs (per-sample ray.put)	140ms (sequential)	~100ms	ThreadPoolExecutor(4) parallelizes puts
Total actor blocking per resim	~430ms	~100ms	~4x less training interruption

Broadcast conditions now stay as shape (1, ...) through _resolve_refs → _chunked_sample, expanded only at the final consumer (_simulate via np.broadcast_to, or embeddings via .expand()).

Background data fetch with tensor pre-building

CachedDataLoader.sync() was blocking training for ~4-5s per epoch due to:

1. checkout_refs + ray.get on the main thread (moved to background in prior commit)
2. _apply_fetch doing per-sample np.array → torch.as_tensor → .to(device) for 1024 samples × N keys

Fix: The background thread now does the full pipeline — checkout_refs, ray.get, np.stack, and torch.as_tensor().to(device) — returning ready-to-use tensors. _apply_fetch on the main thread only does fast batched indexed scatter and torch.cat.

Metric	Before	After	Improvement
_apply_fetch (1024 samples)	4.0s	0.2s	20x
_apply_fetch (256 val samples)	1.2s	0.1s	12x
sync() total	5.3s (blocking)	0.3s (non-blocking)	18x
Epoch time	14.2s	~10s	~30% faster

sync() is now truly non-blocking: checks .done() before applying, skips if the background fetch is still running.

Test results

All smoke tests pass (pytest tests/test_examples_smoke.py):

• ✅ 01_minimal: 2 epochs
• ✅ 02_bimodal: 2 epochs
• ✅ 03_composite: 2 epochs
• ✅ 04_gaussian: 2 epochs
• ✅ 05_linear_regression: 2 epochs

Total pytest runtime: ~3 minutes (CPU mode)

Test plan

• Run pytest tests/test_examples_smoke.py — all 5 passed
• Profile GPU utilization on 05_linear_regression — confirmed 30-70% steady with brief <400ms dips

[cweniger]

Ray Warning Suppression Fix

The "Using blocking ray.get inside async actor" warning has been successfully suppressed.

Root cause: Ray emits this warning via logging.warning(), not warnings.warn(). Previous attempts using warnings.filterwarnings() did not work.

Solution: Set Ray internal blocking_get_inside_async_warned flag to True in NodeWrapper.__init__ before any ray.get calls. This prevents the warning from being logged while preserving the underlying behavior.

The fix is minimal and isolated, with a TODO note about potentially splitting async training actors from synchronous sampling actors in the future.

[cweniger]

Looks all good!