Test streaming

README.md

@@ -18,6 +18,7 @@ cd deepdrivewe
 conda create -n deepdrivewe python=3.10 -y
 conda install omnia::ambertools -y
 conda install conda-forge::openmm==7.7 -y
+conda install anaconda::redis -y
 pip install -U pip setuptools wheel
 pip install -e .
 ```
@@ -33,6 +34,7 @@ ml gcc/14.2.0 cuda/12.5 hdf5
 conda create -n deepdrivewe python=3.12 -y
 conda activate deepdrivewe
 conda install conda-forge::openmm -y
+conda install anaconda::redis -y
 pip install torch --index-url https://download.pytorch.org/whl/cu124
 
 git clone git@github.com:braceal/deepdrivewe.git
@@ -120,6 +122,49 @@ OPENMM_CPU_THREADS=1 nohup python -m deepdrivewe.examples.openmm_ntl9_hk.main --
 Note that we set `OPENMM_CPU_THREADS=1` to restrict each OpenMM simulation to a single thread. This is necessary to prevent
 the simulations from using all available CPU resources. You can also run the simulations on a GPU by adjusting the Parsl configuration.
 
+### Running with streaming
+
+For a full example, see `examples/openmm_ntl9_ddwe_stream`.
+
+To run with streaming, add the following to the config:
+```yaml
+stream_config:
+  # A redis server is used as the stream message broker
+  redis_host: localhost
+  redis_port: 6379
+  # The Store used for stream items is configurable
+  store_config:
+    name: stream-store
+    # Use the same redis server for object storage
+    connector:
+      kind: redis
+      options:
+        hostname: localhost
+        port: 6379
+    # FileConnector example
+    # connector:
+    #   kind: file
+    #   options:
+    #     store_dir: /tmp/proxystore-cache
+```
+
+Then start a redis server in the background:
+```bash
+redis-server --port 6379 --save "" --appendonly no --protected-mode no &> redis.log &
+```
+
+The redis server can later be killed using the job number:
+```bash
+jobs
+kill %<num>
+```
+
+To check resource utilization of the redis server, run the following command:
+```bash
+watch "ps -p \$(pgrep -x redis-server | head -n1) -o pid,comm,%mem,rss,vsz"
+```
+**Note:** If there are multiple redis servers running, this command will choose the first one.
+
 ## Contributing
 
 For development, it is recommended to use a virtual environment. The following

deepdrivewe/examples/openmm_ntl9_ddwe/train.py

@@ -125,9 +125,11 @@ def run_stream_train(
 
     # Loop indefinitely until we get a stop iteration from the stream consumer
     for idx in itertools.count():
+        print(f'Train iteration: {idx}', flush=True)
         # If we have reached the retrain interval, re-initialize the trainer
         # NOTE: This always happens on the first iteration
         if idx % config.stream_retrain_interval == 0:
+            print(f'Retraining model at iteration: {idx}', flush=True)
             # Load the model configuration
             model_config = ConvolutionalVAEConfig.from_yaml(config.config_path)
 
@@ -147,8 +149,17 @@ def run_stream_train(
                 for _ in range(config.stream_items_per_train)
             ]
         except StopIteration:
+            print(
+                f'Reached end of training stream consumer at iteration: {idx}',
+                flush=True,
+            )
             break
 
+        print(
+            f'Got {len(items)} items from stream consumer at iteration: {idx}',
+            flush=True,
+        )
+
         # Extract the contact maps and rmsd from each simulation
         cmaps = np.array([x['contact_maps'] for x in items], dtype=object)
         pcoords = np.array([x['pcoords'] for x in items]).flatten()
@@ -157,11 +168,14 @@ def run_stream_train(
         model_dir = output_dir / f'model_{idx:06d}'
 
         # Fit the model
+        print(f'Fitting model at iteration: {idx}', flush=True)
         checkpoint_path = model.fit(
             x=cmaps,
             model_dir=model_dir,
             scalars={'pcoord': pcoords},
         )
+        print(f'Finished fitting model at iteration: {idx}')
+        print(f'Checkpoint path: {checkpoint_path}')
 
         # Construct the train result
         result = TrainResult(
@@ -170,7 +184,11 @@ def run_stream_train(
         )
 
         # Send the new model weights to the thinker
-        stream_producer.send(topic=TRAIN_TOPIC, obj=result)
+        print(
+            f'Sending new model weights to thinker at iteration: {idx}',
+            flush=True,
+        )
+        stream_producer.send(topic=TRAIN_TOPIC, obj=result, evict=False)
 
     # NOTE: This final return is not necessary, but it is included
     #       to keep the function signature consistent with the non-streaming.

deepdrivewe/simulation/openmm.py

@@ -423,6 +423,9 @@ def report(self, simulation: app.Simulation, state: openmm.State) -> None:
         positions = self.get_positions(simulation, state)
 
         # Collect data from the simulation
+        # NOTE: The individual collectors are responsible for caching
+        # the data they collect, we add it to the data dictionary here
+        # for streaming.
         data = {x.topic: x.collect(positions) for x in self.collectors}
 
         # Stream the data if a stream config is provided

deepdrivewe/workflows/ddwe.py

@@ -348,17 +348,24 @@ def train_stream_processor(self) -> None:
 
             # Clean up the previous training output from the store
             if self.train_output is not None:
+                self.logger.info('Evicting previous training output')
                 # Get the proxy key for the current training output
                 key = get_key(self.train_output)
                 # Evict the key from the store to clean up memory
                 self.stream_config.get_store().evict(key)
+                self.logger.info(
+                    f'Evicted previous training output with key: {key}',
+                )
 
             # Store the training output
             self.train_output = result
 
             # Increment the training iteration
             self.train_iteration += 1
 
+            # Log the training iteration
+            self.logger.info(f'Training iteration: {self.train_iteration}')
+
     def stop_workflow(self) -> None:
         """Stop the workflow."""
         # Set the done flag to signal the agents to stop
@@ -383,7 +390,7 @@ def process_inference_result(self, result: Result) -> None:
 
         # Check if the task failed
         if not result.success:
-            self.logger.warning('Inference failed, quitting workflow.')
+            self.logger.error('Inference failed, quitting workflow.')
             self.stop_workflow()
             return
 

deepdrivewe/workflows/stream.py

@@ -4,6 +4,7 @@
 
 from typing import Any
 
+from parsl.addresses import address_by_hostname
 from proxystore.store import get_store
 from proxystore.store import register_store
 from proxystore.store import Store
@@ -12,6 +13,7 @@
 from proxystore.stream import StreamProducer
 from proxystore.stream.shims.redis import RedisQueuePublisher
 from proxystore.stream.shims.redis import RedisQueueSubscriber
+from pydantic import field_validator
 
 from deepdrivewe import BaseModel
 
@@ -26,6 +28,34 @@ class ProxyStreamConfig(BaseModel):
     redis_host: str = 'localhost'
     redis_port: int = 6379
 
+    @field_validator('redis_host')
+    @classmethod
+    def validate_redis_host(cls, value: str) -> str:
+        """Validate the Redis host."""
+        # Get the hostname if the address is 'hostname'
+        if value == 'hostname':
+            value = address_by_hostname()
+
+        return value
+
+    @field_validator('store_config')
+    @classmethod
+    def validate_store_config(cls, value: StoreConfig) -> StoreConfig:
+        """Validate the store configuration."""
+        if value.connector.kind == 'redis':
+            hostname = value.connector.options.get('hostname')
+            if hostname is None:
+                raise ValueError(
+                    'Hostname is required for Redis connector '
+                    'in store configuration. Use "hostname" to use the '
+                    'hostname of the current machine.',
+                )
+            # If the hostname is 'hostname', look up the hostname and set it
+            if hostname == 'hostname':
+                hostname = address_by_hostname()
+                value.connector.options['hostname'] = hostname
+        return value
+
     def get_store(self) -> Store[Any]:
         """Get the store for the proxy stream.
 
@@ -80,4 +110,4 @@ def get_producer(self, topic: str) -> StreamProducer[Any]:
         """
         store = self.get_store()
         publisher = RedisQueuePublisher(self.redis_host, self.redis_port)
-        return StreamProducer(publisher, {topic: store})
+        return StreamProducer(publisher, stores={topic: store})