Merge pull request #801 from dimensionalOS/rpc-fixes-merge

- implements exception passing to RPC clients (if there is an exception on the server, it will be raised on the caller side)
- stress tests for all transport protocols (we expect 10k messages < sec, 100% reliability expected across all transports)
- stress tests for RPC (1k rpc calls < sec, 100% reliability expected) - SHM actually fails this - has 25% reliability on RPC tests due to how it handles message reception, this is an issue that means we can't use SHM across the board in dimos

- moves rpc call threadpool from LCMRPC to the rpc base so it's used cross-protocol
- shares a single cpp LCM instance/loop per process for all instantiated transports

- investigated kernel networking stats and LCM cpp implementation, optimized details around this to get LCM to perform much better

Author: paul-nechifor

dimos/agents2/spec.py

@@ -162,6 +162,9 @@ def start(self) -> None:
         super().start()
 
     def stop(self) -> None:
+        if hasattr(self, "transport") and self.transport:
+            self.transport.stop()  # type: ignore[attr-defined]
+            self.transport = None  # type: ignore[assignment]
         super().stop()
 
     @rpc

dimos/core/__init__.py

@@ -19,7 +19,7 @@
     pLCMTransport,
     pSHMTransport,
 )
-from dimos.protocol.rpc.lcmrpc import LCMRPC
+from dimos.protocol.rpc import LCMRPC
 from dimos.protocol.rpc.spec import RPCSpec
 from dimos.protocol.tf import LCMTF, TF, PubSubTF, TFConfig, TFSpec
 from dimos.utils.actor_registry import ActorRegistry

dimos/core/module.py

@@ -33,7 +33,7 @@
 from dimos.core.resource import Resource
 from dimos.core.rpc_client import RpcCall
 from dimos.core.stream import In, Out, RemoteIn, RemoteOut, Transport
-from dimos.protocol.rpc import LCMRPC, RPCSpec  # type: ignore[attr-defined]
+from dimos.protocol.rpc import LCMRPC, RPCSpec
 from dimos.protocol.service import Configurable  # type: ignore[attr-defined]
 from dimos.protocol.skill.skill import SkillContainer
 from dimos.protocol.tf import LCMTF, TFSpec

dimos/core/rpc_client.py

@@ -15,7 +15,7 @@
 from collections.abc import Callable
 from typing import Any
 
-from dimos.protocol.rpc.lcmrpc import LCMRPC
+from dimos.protocol.rpc import LCMRPC
 from dimos.utils.logging_config import setup_logger
 
 logger = setup_logger()

dimos/protocol/pubsub/lcmpubsub.py

@@ -64,7 +64,6 @@ class LCMPubSubBase(LCMService, PubSub[Topic, Any]):
     _callbacks: dict[str, list[Callable[[Any], None]]]
 
     def __init__(self, **kwargs) -> None:  # type: ignore[no-untyped-def]
-        LCMService.__init__(self, **kwargs)
         super().__init__(**kwargs)
         self._callbacks = {}
 
@@ -73,6 +72,7 @@ def publish(self, topic: Topic, message: bytes) -> None:
         if self.l is None:
             logger.error("Tried to publish after LCM was closed")
             return
+
         self.l.publish(str(topic), message)
 
     def subscribe(
@@ -88,6 +88,9 @@ def noop() -> None:
 
         lcm_subscription = self.l.subscribe(str(topic), lambda _, msg: callback(msg, topic))
 
+        # Set queue capacity to 10000 to handle high-volume bursts
+        lcm_subscription.set_queue_capacity(10000)
+
         def unsubscribe() -> None:
             if self.l is None:
                 return

dimos/protocol/pubsub/test_spec.py

@@ -23,6 +23,7 @@
 import pytest
 
 from dimos.msgs.geometry_msgs import Vector3
+from dimos.protocol.pubsub.lcmpubsub import LCM, Topic
 from dimos.protocol.pubsub.memory import Memory
 
 
@@ -61,27 +62,21 @@ def redis_context():
     print("Redis not available")
 
 
-try:
-    from dimos.protocol.pubsub.lcmpubsub import LCM, Topic
+@contextmanager
+def lcm_context():
+    lcm_pubsub = LCM(autoconf=True)
+    lcm_pubsub.start()
+    yield lcm_pubsub
+    lcm_pubsub.stop()
 
-    @contextmanager
-    def lcm_context():
-        lcm_pubsub = LCM(autoconf=True)
-        lcm_pubsub.start()
-        yield lcm_pubsub
-        lcm_pubsub.stop()
 
-    testdata.append(
-        (
-            lcm_context,
-            Topic(topic="/test_topic", lcm_type=Vector3),
-            [Vector3(1, 2, 3), Vector3(4, 5, 6), Vector3(7, 8, 9)],  # Using Vector3 as mock data,
-        )
+testdata.append(
+    (
+        lcm_context,
+        Topic(topic="/test_topic", lcm_type=Vector3),
+        [Vector3(1, 2, 3), Vector3(4, 5, 6), Vector3(7, 8, 9)],  # Using Vector3 as mock data,
     )
-
-except (ConnectionError, ImportError):
-    # either redis is not installed or the server is not running
-    print("LCM not available")
+)
 
 
 from dimos.protocol.pubsub.shmpubsub import PickleSharedMemory
@@ -263,3 +258,40 @@ async def consume_messages() -> None:
         # Verify all messages were received in order
         assert len(received_messages) == len(messages_to_send)
         assert received_messages == messages_to_send
+
+
+@pytest.mark.parametrize("pubsub_context, topic, values", testdata)
+def test_high_volume_messages(pubsub_context, topic, values) -> None:
+    """Test that all 5000 messages are received correctly."""
+    with pubsub_context() as x:
+        # Create a list to capture received messages
+        received_messages = []
+        last_message_time = [time.time()]  # Use list to allow modification in callback
+
+        # Define callback function
+        def callback(message, topic) -> None:
+            received_messages.append(message)
+            last_message_time[0] = time.time()
+
+        # Subscribe to the topic
+        x.subscribe(topic, callback)
+
+        # Publish 10000 messages
+        num_messages = 10000
+        for _ in range(num_messages):
+            x.publish(topic, values[0])
+
+        # Wait until no messages received for 0.5 seconds
+        timeout = 1.0  # Maximum time to wait
+        stable_duration = 0.1  # Time without new messages to consider done
+        start_time = time.time()
+
+        while time.time() - start_time < timeout:
+            if time.time() - last_message_time[0] >= stable_duration:
+                break
+            time.sleep(0.1)
+
+        # Capture count and clear list to avoid printing huge list on failure
+        received_len = len(received_messages)
+        received_messages.clear()
+        assert received_len == num_messages, f"Expected {num_messages} messages, got {received_len}"

dimos/protocol/rpc/__init__.py

@@ -12,5 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from dimos.protocol.rpc.lcmrpc import LCMRPC
+from dimos.protocol.rpc.pubsubrpc import LCMRPC, ShmRPC
 from dimos.protocol.rpc.spec import RPCClient, RPCServer, RPCSpec
+
+__all__ = ["LCMRPC", "RPCClient", "RPCServer", "RPCSpec", "ShmRPC"]