Improve testing harness using psutil process handling

container/deps/requirements.txt

@@ -17,6 +17,7 @@ fastapi==0.115.6
 ftfy
 grpcio-tools==1.66.0
 httpx
+msgspec # Event plane should have this in valid wheel configuration
 mypy
 numpy
 opentelemetry-api

icp/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

icp/examples/python/event_plane/README.md

@@ -0,0 +1,55 @@
+<!--
+SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+SPDX-License-Identifier: Apache-2.0
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+-->
+
+# Event Plane example
+
+A basic example that demonstrates how to use the Event Plane API to create an event plane, register an event, and trigger the event.
+
+## Code overview
+
+### 1) Initialize NATS server and create an event plane
+```python
+    server_url = "nats://localhost:4222"
+    component_id = uuid.uuid4()
+    plane = NatsEventPlane(server_url, component_id)
+    await plane.connect()
+```
+
+### 2) Define the callback function for receiving events
+```python
+    received_events = []
+    async def callback(event, metadata):
+        print(metadata)
+        received_events.append(metadata)
+```
+
+### 3) Prepare the event event_topic, event type, and event payload
+```python
+    event_topic = EventTopic(["test", "event_topic"])
+    event_type = "test_event"
+    event = b"my_payload"
+```
+
+### 4) Subscribe to the event event_topic and type and register the callback function
+```python
+    await plane.subscribe(callback, event_topic=event_topic, event_type=event_type)
+```
+
+### 5) Publish the event
+```python
+    await plane.publish(event, event_type, event_topic)
+```

icp/examples/python/event_plane/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

icp/examples/python/event_plane/parser.py

@@ -0,0 +1,47 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+
+from triton_distributed.icp.nats_event_plane import DEFAULT_EVENTS_PORT
+
+
+def parse_args(args=None):
+    parser = argparse.ArgumentParser(description="Event Plane Example")
+
+    parser.add_argument(
+        "--nats-port",
+        type=int,
+        default=DEFAULT_EVENTS_PORT,
+        help="Nats server port",
+    )
+
+    parser.add_argument(
+        "--publisher-count",
+        type=int,
+        default=1,
+        help="Number of publishers to deploy",
+    )
+
+    parser.add_argument(
+        "--subscriber-count",
+        type=int,
+        default=10,
+        help="Number of subscribers to deploy",
+    )
+
+    args = parser.parse_args(args)
+
+    return args

icp/examples/python/event_plane/pub_sub_simple.py

@@ -0,0 +1,60 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import asyncio
+import uuid
+
+from triton_distributed.icp.nats_event_plane import (
+    DEFAULT_EVENTS_URI,
+    EventTopic,
+    NatsEventPlane,
+)
+
+
+async def single_publisher_subscriber_example():
+    # async with aclosing(event_plane()) as event_plane_instance:
+    # event_plane_instance = await anext(event_plane)
+
+    server_url = DEFAULT_EVENTS_URI
+    component_id = str(uuid.uuid4())
+    plane = NatsEventPlane(server_url, component_id)
+
+    await plane.connect()
+    received_events = []
+
+    async def callback(event):
+        print(event)
+        received_events.append(event)
+
+    event_topic = EventTopic(["test", "event_topic"])
+    event_type = "test_event"
+    event = b"my_payload"
+
+    await plane.subscribe(callback, event_topic=event_topic, event_type=event_type)
+
+    await plane.publish(event, event_type, event_topic)
+
+    # Allow time for message to propagate
+    await asyncio.sleep(3)
+
+    print(f"received_events: {received_events}")
+    # assert received_events[0][0].event_id == event.event_id
+
+    await plane.disconnect()
+
+
+if __name__ == "__main__":
+    asyncio.run(single_publisher_subscriber_example())

icp/examples/python/event_plane/publisher.py

@@ -0,0 +1,91 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import argparse
+import asyncio
+import logging
+import json
+import uuid
+
+from triton_distributed.icp.nats_event_plane import (
+    DEFAULT_EVENTS_URI,
+    EventTopic,
+    NatsEventPlane,
+)
+
+logging.basicConfig(
+    level=logging.DEBUG,
+    format="%(filename)s: %(levelname)s: %(funcName)s(): %(lineno)d:\t%(message)s",
+)
+
+logger = logging.getLogger(__name__)
+
+async def main(args):
+    server_url = DEFAULT_EVENTS_URI
+    event_plane = NatsEventPlane(server_url, args.component_id)
+
+    await event_plane.connect()
+
+    try:
+        event_topic = EventTopic(["publisher", str(args.publisher_id)])
+
+        for i in range(args.event_count):
+            event = f"Payload from publisher {args.publisher_id} idx {i}".encode()
+            await event_plane.publish(event, 
+                                      args.event_type, 
+                                      event_topic)
+            logger.info(f"Published event from publisher {args.publisher_id}")
+            # Serialize sent events to JSON
+            if args.save_events_path:
+                with open(args.save_events_path, "a+") as json_file:
+                    event_obj = {"event_payload": str(event.decode("utf-8")),
+                            "event_id": str(args.publisher_id), 
+                            "event_topic": str(event_topic),
+                            "event_type": args.event_type,
+                            "component_id": str(args.component_id)}
+                    json_file.write(json.dumps(event_obj))
+                    json_file.write("\n")
+            await asyncio.sleep(0.01)
+    finally:
+        await event_plane.disconnect()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Event publisher script")
+    parser.add_argument(
+        "--component-id",
+        type=uuid.UUID,
+        default=uuid.uuid4(),
+        help="Component ID (UUID)",
+    )
+    parser.add_argument(
+        "--event-type", type=str, default="test_event", help="Event type"
+    )
+    parser.add_argument("--publisher-id", type=int, required=True, help="Publisher ID")
+    parser.add_argument(
+        "--event-count", type=int, default=10, help="Event count to be published."
+    )
+    parser.add_argument(
+        "--save-events-path",
+        type=str,
+        default=None,
+        help="Path to save received events as JSON",
+    )
+
+    args = parser.parse_args()
+    asyncio.run(
+        main(args)
+    )