refactor: Simplify networking

81cf6bce64f9de306fe906cc237dcf579333b002 · 2025-07-07 19:32:21 +0100 · Arbion Halili
Files touched

M shared/types/networking/data_plane.py
M shared/types/networking/services.py
M shared/types/networking/topology.py
Diff

commit 81cf6bce64f9de306fe906cc237dcf579333b002
Author: Arbion Halili <99731180+ToxicPine@users.noreply.github.com>
Date:   Mon Jul 7 19:32:21 2025 +0100

    refactor: Simplify networking
---
 shared/types/networking/data_plane.py | 73 ++++++++++++++---------------------
 shared/types/networking/services.py   | 20 +++-------
 shared/types/networking/topology.py   | 12 +++---
 3 files changed, 39 insertions(+), 66 deletions(-)

diff --git a/shared/types/networking/data_plane.py b/shared/types/networking/data_plane.py
index 7607a1c2..acb022eb 100644
--- a/shared/types/networking/data_plane.py
+++ b/shared/types/networking/data_plane.py
@@ -1,13 +1,9 @@
 from enum import Enum
-from typing import Generic, Mapping, Tuple, TypeVar, final
+from typing import Annotated, Literal, TypeVar, Union, final
 
-from pydantic import BaseModel, IPvAnyAddress
+from pydantic import BaseModel, Field, IPvAnyAddress, TypeAdapter
 
-from shared.types.common import NewUUID, NodeId
-from shared.types.graphs.common import (
-    Edge,
-    EdgeData,
-)
+from shared.types.common import NewUUID
 
 
 class DataPlaneEdgeId(NewUUID):
@@ -15,7 +11,7 @@ class DataPlaneEdgeId(NewUUID):
 
 
 class AddressingProtocol(str, Enum):
-    IPvAny = "IPvAny"
+    IPvAnyAddress = "IPvAnyAddress"
 
 
 class ApplicationProtocol(str, Enum):
@@ -27,56 +23,43 @@ ApP = TypeVar("ApP", bound=ApplicationProtocol)
 
 
 @final
-class EdgeDataTransferRate(BaseModel):
+class DataPlaneEdgeBenchmarkData(BaseModel):
     throughput: float
     latency: float
     jitter: float
 
 
-class DataPlaneEdgeMetadata(BaseModel, Generic[AdP, ApP]): ...
+class CommonDataPlaneEdgeData(BaseModel):
+    edge_data_transfer_rate: DataPlaneEdgeBenchmarkData | None = None
 
 
-@final
-class DataPlaneEdgeType(BaseModel, Generic[AdP, ApP]):
-    addressing_protocol: AdP
-    application_protocol: ApP
-
-
-@final
-class MLXEdgeContext(
-    DataPlaneEdgeMetadata[AddressingProtocol.IPvAny, ApplicationProtocol.MLX]
-):
+class MlxEdgeMetadata(BaseModel):
     source_ip: IPvAnyAddress
     sink_ip: IPvAnyAddress
 
 
-class DataPlaneEdgeInfoType(str, Enum):
-    network_profile = "network_profile"
-    other = "other"
-
-
-AllDataPlaneEdgeInfo = Tuple[DataPlaneEdgeInfoType.network_profile]
-
-
-DataPlaneEdgeInfoTypeT = TypeVar(
-    "DataPlaneEdgeInfoTypeT", bound=DataPlaneEdgeInfoType, covariant=True
-)
-
-
-class DataPlaneEdgeInfo(BaseModel, Generic[DataPlaneEdgeInfoTypeT]):
-    edge_info_type: DataPlaneEdgeInfoTypeT
-
-
-SetOfEdgeInfo = TypeVar("SetOfEdgeInfo", bound=Tuple[DataPlaneEdgeInfoType, ...])
-
+class BaseDataPlaneEdgeData[AdP: AddressingProtocol, ApP: ApplicationProtocol](
+    BaseModel
+):
+    addressing_protocol: AdP
+    application_protocol: ApP
+    common_data: CommonDataPlaneEdgeData
 
-class DataPlaneEdgeData(EdgeData[DataPlaneEdgeType[AdP, ApP]], Generic[AdP, ApP]):
-    edge_info: Mapping[DataPlaneEdgeInfoType, DataPlaneEdgeInfo[DataPlaneEdgeInfoType]]
-    edge_metadata: DataPlaneEdgeMetadata[AdP, ApP]
 
+class MlxEdge(
+    BaseDataPlaneEdgeData[AddressingProtocol.IPvAnyAddress, ApplicationProtocol.MLX]
+):
+    addressing_protocol: Literal[AddressingProtocol.IPvAnyAddress] = (
+        AddressingProtocol.IPvAnyAddress
+    )
+    application_protocol: Literal[ApplicationProtocol.MLX] = ApplicationProtocol.MLX
+    mlx_metadata: MlxEdgeMetadata
 
-class DataPlaneEdgeProfile(DataPlaneEdgeInfo[DataPlaneEdgeInfoTypeT]):
-    edge_data_transfer_rate: EdgeDataTransferRate
 
+DataPlaneEdgeData = Union[MlxEdge]
 
-class DataPlaneEdge(Edge[DataPlaneEdgeType[AdP, ApP], DataPlaneEdgeId, NodeId]): ...
+_DataPlaneEdgeData = Annotated[
+    DataPlaneEdgeData,
+    Field(discriminator="addressing_protocol"),
+]
+DataPlaneEdgeAdapter: TypeAdapter[DataPlaneEdgeData] = TypeAdapter(_DataPlaneEdgeData)
diff --git a/shared/types/networking/services.py b/shared/types/networking/services.py
index 51620421..01655d15 100644
--- a/shared/types/networking/services.py
+++ b/shared/types/networking/services.py
@@ -1,4 +1,4 @@
-from typing import Callable, NewType, Protocol, TypeVar
+from typing import Callable, NewType, Protocol
 
 from shared.types.networking.control_plane import (
     ControlPlaneEdgeId,
@@ -7,10 +7,7 @@ from shared.types.networking.control_plane import (
 
 TopicName = NewType("TopicName", str)
 
-MessageT = TypeVar("MessageT", bound=object)
-
-
-PubSubMessageHandler = Callable[[TopicName, MessageT], None]
+PubSubMessageHandler = Callable[[TopicName, object], None]
 NodeConnectedHandler = Callable[
     [
         ControlPlaneEdgeId,
@@ -22,16 +19,11 @@ NodeDisconnectedHandler = Callable[[ControlPlaneEdgeId], None]
 
 
 class DiscoveryService(Protocol):
-    def register_node_connected_handler(
-        self, handler: NodeConnectedHandler
-    ) -> None: ...
-    def register_node_disconnected_handler(
-        self, handler: NodeDisconnectedHandler
-    ) -> None: ...
+    def on_node_connected(self, handler: NodeConnectedHandler) -> None: ...
+    def on_node_disconnected(self, handler: NodeDisconnectedHandler) -> None: ...
 
 
 class PubSubService(Protocol):
-    def register_handler(
-        self, key: str, topic_name: TopicName, handler: PubSubMessageHandler[MessageT]
+    def on_message_received(
+        self, topic_name: TopicName, handler: PubSubMessageHandler
     ) -> None: ...
-    def deregister_handler(self, key: str) -> None: ...
diff --git a/shared/types/networking/topology.py b/shared/types/networking/topology.py
index f59d7064..61e8900b 100644
--- a/shared/types/networking/topology.py
+++ b/shared/types/networking/topology.py
@@ -2,9 +2,7 @@ from shared.types.common import NodeId
 from shared.types.graphs.common import Graph, GraphData
 from shared.types.networking.control_plane import ControlPlaneEdgeId
 from shared.types.networking.data_plane import (
-    AddressingProtocol,
-    ApplicationProtocol,
-    DataPlaneEdge,
+    DataPlaneEdgeData,
     DataPlaneEdgeId,
 )
 from shared.types.worker.common import NodeStatus
@@ -12,14 +10,14 @@ from shared.types.worker.common import NodeStatus
 
 class DataPlaneTopology(
     Graph[
-        DataPlaneEdge[AddressingProtocol, ApplicationProtocol],
+        DataPlaneEdgeData,
         None,
         DataPlaneEdgeId,
         NodeId,
     ]
 ):
     graph_data: GraphData[
-        DataPlaneEdge[AddressingProtocol, ApplicationProtocol],
+        DataPlaneEdgeData,
         None,
         DataPlaneEdgeId,
         NodeId,
@@ -28,14 +26,14 @@ class DataPlaneTopology(
 
 class OrphanedPartOfDataPlaneTopology(
     Graph[
-        DataPlaneEdge[AddressingProtocol, ApplicationProtocol],
+        DataPlaneEdgeData,
         None,
         DataPlaneEdgeId,
         NodeId,
     ]
 ):
     graph_data: GraphData[
-        DataPlaneEdge[AddressingProtocol, ApplicationProtocol],
+        DataPlaneEdgeData,
         None,
         DataPlaneEdgeId,
         NodeId,
← 6c8b8b30 added rust to flake · back to Exo · refactor: A Lot e1894bc1 →