Skip to content

Test streaming - #40

Open
braceal wants to merge 26 commits into
mainfrom
test-streaming
Open

Test streaming#40
braceal wants to merge 26 commits into
mainfrom
test-streaming

Conversation

@braceal

@braceal braceal commented Oct 30, 2025

Copy link
Copy Markdown
Collaborator

Gets the streaming DDWE workflow (DDWEStreamThinker) running end-to-end on Vista, and adds a complete streaming example plus setup docs.

Four bugs found while debugging this branch were split out and merged separately: #48 (WeightedEnsemble.metadata default), #49 (inverted use_stale_model gating), #50 (print_errors on null failure_info), #51 (unbounded re-train history). main has been merged back in, so what remains here is only the streaming feature work.

Stream configuration

ProxyStreamConfig gains two validators (workflows/stream.py):

  • redis_host accepts the sentinel value hostname, resolved to address_by_hostname() at load time. This is what lets one config file work unchanged across compute nodes, where the broker address is not known until the job lands.
  • store_config requires a hostname when the connector kind is redis, and accepts the same hostname sentinel. Previously a missing hostname surfaced later as an opaque connector error; it now fails at config load with a message that names the fix.

StreamProducer(publisher, stores={topic: store}) now passes stores as a keyword, per proxystore/proxystore#641.

proxystore is installed as proxystore[redis].

Proxy lifetime for streamed model weights

The training producer sends its TrainResult with evict=False. The default eviction-on-consume is wrong here: the thinker holds the model across iterations and hands it to every inference task until a newer one arrives, so evicting at consume time pulls the object out from under in-flight inference.

Memory is still reclaimed — train_stream_processor explicitly evicts the previous proxy's key when it is replaced, which is now logged on both sides of the call.

Example and docs

  • New examples/openmm_ntl9_ddwe_stream/ — config, CVAE config, NTL9 basis state, reference structure, and a pretrained checkpoint for a runnable streaming NTL9 DDWE run.
  • README gains a Running with streaming section: the stream_config block with a commented FileConnector alternative, how to start and kill a background redis server, and how to watch its memory usage. The install instructions add conda install anaconda::redis.

Notes for reviewers

  • Much of the remaining diff is print / logger.info instrumentation added while debugging the stream path. Worth deciding what should stay as permanent logging versus what should be dropped before merge.
  • examples/openmm_ntl9_ddwe_stream/checkpoint-epoch-100.pt is a 13 MB binary committed directly to the repo. Flagging in case it belongs somewhere out-of-tree.
  • Inference failure now logs at error rather than warning, since it stops the workflow.
  • The commit history is debug iteration (debug, logs, config, …) and would be better squashed on merge.

Base automatically changed from vista-stream to main August 11, 2026 21:10
# Conflicts:
#	deepdrivewe/examples/openmm_ntl9_ddwe/train.py
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