← back to Exo
refactor: Distinguish the topology of the control plane from that of the data plane
10224d09ded0f332fd023da4b8682b0da049a6ba · 2025-07-03 15:45:54 +0100 · Arbion Halili
Files touched
M master/api.pyM shared/types/events/common.pyM shared/types/events/events.pyM shared/types/graphs/resource_graph.pyA shared/types/networking/control_plane.pyA shared/types/networking/data_plane.pyD shared/types/networking/edges.pyM shared/types/networking/services.pyM shared/types/networking/topology.pyM shared/types/profiling/common.pyM shared/types/states/master.pyM shared/types/states/shared.pyM shared/types/states/worker.pyM shared/types/tasks/common.py
Diff
commit 10224d09ded0f332fd023da4b8682b0da049a6ba
Author: Arbion Halili <99731180+ToxicPine@users.noreply.github.com>
Date: Thu Jul 3 15:45:54 2025 +0100
refactor: Distinguish the topology of the control plane from that of the data plane
---
master/api.py | 6 +-
shared/types/events/common.py | 25 +++++----
shared/types/events/events.py | 94 +++++++++++++++++++++++++-------
shared/types/graphs/resource_graph.py | 11 ++--
shared/types/networking/control_plane.py | 11 ++++
shared/types/networking/data_plane.py | 82 ++++++++++++++++++++++++++++
shared/types/networking/edges.py | 82 ----------------------------
shared/types/networking/services.py | 17 ++----
shared/types/networking/topology.py | 60 +++++++++++++++-----
shared/types/profiling/common.py | 2 +-
shared/types/states/master.py | 51 ++++++++++-------
shared/types/states/shared.py | 10 +++-
shared/types/states/worker.py | 4 +-
shared/types/tasks/common.py | 26 +++++----
14 files changed, 299 insertions(+), 182 deletions(-)
diff --git a/master/api.py b/master/api.py
index cc45c786..28c35ce1 100644
--- a/master/api.py
+++ b/master/api.py
@@ -2,14 +2,16 @@ from typing import Protocol
from shared.types.models.common import Model, ModelId
from shared.types.models.sources import ModelSource
-from shared.types.networking.topology import Topology
+from shared.types.networking.topology import ControlPlaneTopology, DataPlaneTopology
from shared.types.worker.common import InstanceId
from shared.types.worker.downloads import DownloadProgress
from shared.types.worker.instances import Instance
class ControlPlaneAPI(Protocol):
- def get_topology(self) -> Topology: ...
+ def get_control_plane_topology(self) -> ControlPlaneTopology: ...
+
+ def get_data_plane_topology(self) -> DataPlaneTopology: ...
def list_instances(self) -> list[Instance]: ...
diff --git a/shared/types/events/common.py b/shared/types/events/common.py
index 1663fae8..ab920306 100644
--- a/shared/types/events/common.py
+++ b/shared/types/events/common.py
@@ -38,6 +38,7 @@ class StreamingEventTypes(str, Enum):
class InstanceEventTypes(str, Enum):
InstanceCreated = "InstanceCreated"
InstanceDeleted = "InstanceDeleted"
+ InstanceToBeReplacedAtomically = "InstanceToBeReplacedAtomically"
InstanceReplacedAtomically = "InstanceReplacedAtomically"
InstanceStatusUpdated = "InstanceStatusUpdated"
@@ -46,18 +47,20 @@ class InstanceStateEventTypes(str, Enum):
InstanceRunnerStateUpdated = "InstanceRunnerStateUpdated"
-class NodeStatusEventTypes(str, Enum):
- NodeStatusUpdated = "NodeStatusUpdated"
+class NodePerformanceEventTypes(str, Enum):
+ NodePerformanceProfiled = "NodePerformanceProfiled"
-class NodeProfileEventTypes(str, Enum):
- NodeProfileUpdated = "NodeProfileUpdated"
+class DataPlaneEventTypes(str, Enum):
+ DataPlaneEdgeCreated = "DataPlaneEdgeCreated"
+ DataPlaneEdgeProfiled = "DataPlaneEdgeProfiled"
+ DataPlaneEdgeDeleted = "DataPlaneEdgeDeleted"
-class EdgeEventTypes(str, Enum):
- EdgeCreated = "EdgeCreated"
- EdgeUpdated = "EdgeUpdated"
- EdgeDeleted = "EdgeDeleted"
+class ControlPlaneEventTypes(str, Enum):
+ WorkerConnected = "WorkerConnected"
+ WorkerStatusUpdated = "WorkerStatusUpdated"
+ WorkerDisconnected = "WorkerDisconnected"
class TimerEventTypes(str, Enum):
@@ -70,9 +73,9 @@ EventTypes = Union[
StreamingEventTypes,
InstanceEventTypes,
InstanceStateEventTypes,
- NodeStatusEventTypes,
- NodeProfileEventTypes,
- EdgeEventTypes,
+ NodePerformanceEventTypes,
+ ControlPlaneEventTypes,
+ DataPlaneEventTypes,
TimerEventTypes,
MLXEventTypes,
]
diff --git a/shared/types/events/events.py b/shared/types/events/events.py
index 22a6dd89..cd0da509 100644
--- a/shared/types/events/events.py
+++ b/shared/types/events/events.py
@@ -1,22 +1,35 @@
from __future__ import annotations
-from typing import Any, Literal, Tuple
+from typing import Any, Generic, Literal, Tuple, TypeVar
from pydantic import BaseModel
from shared.types.common import NewUUID, NodeId
from shared.types.events.common import (
+ ControlPlaneEventTypes,
+ DataPlaneEventTypes,
Event,
InstanceEventTypes,
InstanceStateEventTypes,
MLXEventTypes,
- NodeProfileEventTypes,
- NodeStatusEventTypes,
+ NodePerformanceEventTypes,
StreamingEventTypes,
TaskEventTypes,
TimerEventTypes,
)
-from shared.types.profiling.common import NodeProfile
+from shared.types.networking.control_plane import (
+ ControlPlaneEdgeId,
+ ControlPlaneEdgeType,
+)
+from shared.types.networking.data_plane import (
+ AddressingProtocol,
+ ApplicationProtocol,
+ DataPlaneEdge,
+ DataPlaneEdgeId,
+ DataPlaneEdgeInfoType,
+ DataPlaneEdgeProfile,
+)
+from shared.types.profiling.common import NodePerformanceProfile
from shared.types.tasks.common import (
TaskData,
TaskId,
@@ -41,18 +54,21 @@ class TimerData(BaseModel):
timer_id: TimerId
-class TaskCreated(Event[TaskEventTypes.TaskCreated]):
+TaskTypeT = TypeVar("TaskTypeT", bound=TaskType)
+
+
+class TaskCreated(Event[TaskEventTypes.TaskCreated], Generic[TaskTypeT]):
event_type: Literal[TaskEventTypes.TaskCreated] = TaskEventTypes.TaskCreated
task_id: TaskId
- task_data: TaskData[TaskType]
- task_state: TaskUpdate[Literal[TaskStatusType.Pending]]
+ task_data: TaskData[TaskTypeT]
+ task_state: TaskUpdate[Literal[TaskStatusType.Pending], TaskTypeT]
on_instance: InstanceId
-class TaskUpdated(Event[TaskEventTypes.TaskUpdated]):
+class TaskUpdated(Event[TaskEventTypes.TaskUpdated], Generic[TaskTypeT]):
event_type: Literal[TaskEventTypes.TaskUpdated] = TaskEventTypes.TaskUpdated
task_id: TaskId
- update_data: TaskUpdate[TaskStatusType]
+ update_data: TaskUpdate[TaskStatusType, TaskTypeT]
class TaskDeleted(Event[TaskEventTypes.TaskDeleted]):
@@ -66,6 +82,7 @@ class InstanceCreated(Event[InstanceEventTypes.InstanceCreated]):
)
instance_id: InstanceId
instance_data: InstanceData
+ target_status: InstanceStatus
class InstanceDeleted(Event[InstanceEventTypes.InstanceDeleted]):
@@ -93,13 +110,17 @@ class InstanceRunnerStateUpdated(
state_update: Tuple[RunnerId, RunnerState[RunnerStateType]]
+class InstanceToBeReplacedAtomically(
+ Event[InstanceEventTypes.InstanceToBeReplacedAtomically]
+):
+ transition: Tuple[InstanceId, InstanceId]
+
+
class InstanceReplacedAtomically(Event[InstanceEventTypes.InstanceReplacedAtomically]):
event_type: Literal[InstanceEventTypes.InstanceReplacedAtomically] = (
InstanceEventTypes.InstanceReplacedAtomically
)
- old_instance_id: InstanceId
- new_instance_id: InstanceId
- new_instance_data: InstanceData
+ transition: Tuple[InstanceId, InstanceId]
class MLXInferenceSagaPrepare(Event[MLXEventTypes.MLXInferenceSagaPrepare]):
@@ -118,22 +139,36 @@ class MLXInferenceSagaStartPrepare(Event[MLXEventTypes.MLXInferenceSagaStartPrep
instance_id: InstanceId
-class NodeProfileUpdated(Event[NodeProfileEventTypes.NodeProfileUpdated]):
- event_type: Literal[NodeProfileEventTypes.NodeProfileUpdated] = (
- NodeProfileEventTypes.NodeProfileUpdated
+class NodePerformanceProfiled(Event[NodePerformanceEventTypes.NodePerformanceProfiled]):
+ event_type: Literal[NodePerformanceEventTypes.NodePerformanceProfiled] = (
+ NodePerformanceEventTypes.NodePerformanceProfiled
)
node_id: NodeId
- node_profile: NodeProfile
+ node_profile: NodePerformanceProfile
-class NodeStatusUpdated(Event[NodeStatusEventTypes.NodeStatusUpdated]):
- event_type: Literal[NodeStatusEventTypes.NodeStatusUpdated] = (
- NodeStatusEventTypes.NodeStatusUpdated
+class WorkerConnected(Event[ControlPlaneEventTypes.WorkerConnected]):
+ event_type: Literal[ControlPlaneEventTypes.WorkerConnected] = (
+ ControlPlaneEventTypes.WorkerConnected
+ )
+ edge: DataPlaneEdge[AddressingProtocol, ApplicationProtocol]
+
+
+class WorkerStatusUpdated(Event[ControlPlaneEventTypes.WorkerStatusUpdated]):
+ event_type: Literal[ControlPlaneEventTypes.WorkerStatusUpdated] = (
+ ControlPlaneEventTypes.WorkerStatusUpdated
)
node_id: NodeId
node_state: NodeStatus
+class WorkerDisconnected(Event[ControlPlaneEventTypes.WorkerConnected]):
+ event_type: Literal[ControlPlaneEventTypes.WorkerConnected] = (
+ ControlPlaneEventTypes.WorkerConnected
+ )
+ vertex_id: ControlPlaneEdgeId
+
+
class ChunkGenerated(Event[StreamingEventTypes.ChunkGenerated]):
event_type: Literal[StreamingEventTypes.ChunkGenerated] = (
StreamingEventTypes.ChunkGenerated
@@ -143,6 +178,27 @@ class ChunkGenerated(Event[StreamingEventTypes.ChunkGenerated]):
chunk: Any
+class DataPlaneEdgeCreated(Event[DataPlaneEventTypes.DataPlaneEdgeCreated]):
+ event_type: Literal[DataPlaneEventTypes.DataPlaneEdgeCreated] = (
+ DataPlaneEventTypes.DataPlaneEdgeCreated
+ )
+ vertex: ControlPlaneEdgeType
+
+
+class DataPlaneEdgeProfiled(Event[DataPlaneEventTypes.DataPlaneEdgeProfiled]):
+ event_type: Literal[DataPlaneEventTypes.DataPlaneEdgeProfiled] = (
+ DataPlaneEventTypes.DataPlaneEdgeProfiled
+ )
+ edge_profile: DataPlaneEdgeProfile[Literal[DataPlaneEdgeInfoType.network_profile]]
+
+
+class DataPlaneEdgeDeleted(Event[DataPlaneEventTypes.DataPlaneEdgeDeleted]):
+ event_type: Literal[DataPlaneEventTypes.DataPlaneEdgeDeleted] = (
+ DataPlaneEventTypes.DataPlaneEdgeDeleted
+ )
+ edge_id: DataPlaneEdgeId
+
+
class TimerScheduled(Event[TimerEventTypes.TimerCreated]):
event_type: Literal[TimerEventTypes.TimerCreated] = TimerEventTypes.TimerCreated
timer_data: TimerData
diff --git a/shared/types/graphs/resource_graph.py b/shared/types/graphs/resource_graph.py
index 4c469d9c..8f664507 100644
--- a/shared/types/graphs/resource_graph.py
+++ b/shared/types/graphs/resource_graph.py
@@ -3,16 +3,15 @@ from collections.abc import Mapping
from pydantic import BaseModel
from shared.types.common import NodeId
-from shared.types.networking.topology import Topology
-from shared.types.profiling.common import NodeProfile
-from shared.types.worker.common import NodeStatus
+from shared.types.networking.topology import ControlPlaneTopology, DataPlaneTopology
+from shared.types.profiling.common import NodePerformanceProfile
class ResourceGraph(BaseModel): ...
def get_graph_of_compute_resources(
- network_topology: Topology,
- node_statuses: Mapping[NodeId, NodeStatus],
- node_profiles: Mapping[NodeId, NodeProfile],
+ control_plane_topology: ControlPlaneTopology,
+ data_plane_topology: DataPlaneTopology,
+ node_profiles: Mapping[NodeId, NodePerformanceProfile],
) -> ResourceGraph: ...
diff --git a/shared/types/networking/control_plane.py b/shared/types/networking/control_plane.py
new file mode 100644
index 00000000..574ff097
--- /dev/null
+++ b/shared/types/networking/control_plane.py
@@ -0,0 +1,11 @@
+from typing import TypeAlias
+
+from shared.types.common import NewUUID, NodeId
+from shared.types.graphs.common import Edge
+
+
+class ControlPlaneEdgeId(NewUUID):
+ pass
+
+
+ControlPlaneEdgeType: TypeAlias = Edge[None, ControlPlaneEdgeId, NodeId]
diff --git a/shared/types/networking/data_plane.py b/shared/types/networking/data_plane.py
new file mode 100644
index 00000000..7607a1c2
--- /dev/null
+++ b/shared/types/networking/data_plane.py
@@ -0,0 +1,82 @@
+from enum import Enum
+from typing import Generic, Mapping, Tuple, TypeVar, final
+
+from pydantic import BaseModel, IPvAnyAddress
+
+from shared.types.common import NewUUID, NodeId
+from shared.types.graphs.common import (
+ Edge,
+ EdgeData,
+)
+
+
+class DataPlaneEdgeId(NewUUID):
+ pass
+
+
+class AddressingProtocol(str, Enum):
+ IPvAny = "IPvAny"
+
+
+class ApplicationProtocol(str, Enum):
+ MLX = "MLX"
+
+
+AdP = TypeVar("AdP", bound=AddressingProtocol)
+ApP = TypeVar("ApP", bound=ApplicationProtocol)
+
+
+@final
+class EdgeDataTransferRate(BaseModel):
+ throughput: float
+ latency: float
+ jitter: float
+
+
+class DataPlaneEdgeMetadata(BaseModel, Generic[AdP, ApP]): ...
+
+
+@final
+class DataPlaneEdgeType(BaseModel, Generic[AdP, ApP]):
+ addressing_protocol: AdP
+ application_protocol: ApP
+
+
+@final
+class MLXEdgeContext(
+ DataPlaneEdgeMetadata[AddressingProtocol.IPvAny, ApplicationProtocol.MLX]
+):
+ 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 DataPlaneEdgeData(EdgeData[DataPlaneEdgeType[AdP, ApP]], Generic[AdP, ApP]):
+ edge_info: Mapping[DataPlaneEdgeInfoType, DataPlaneEdgeInfo[DataPlaneEdgeInfoType]]
+ edge_metadata: DataPlaneEdgeMetadata[AdP, ApP]
+
+
+class DataPlaneEdgeProfile(DataPlaneEdgeInfo[DataPlaneEdgeInfoTypeT]):
+ edge_data_transfer_rate: EdgeDataTransferRate
+
+
+class DataPlaneEdge(Edge[DataPlaneEdgeType[AdP, ApP], DataPlaneEdgeId, NodeId]): ...
diff --git a/shared/types/networking/edges.py b/shared/types/networking/edges.py
deleted file mode 100644
index 3c90a837..00000000
--- a/shared/types/networking/edges.py
+++ /dev/null
@@ -1,82 +0,0 @@
-from enum import Enum
-from typing import Generic, Mapping, Tuple, TypeVar, final
-
-from pydantic import BaseModel, IPvAnyAddress
-
-from shared.types.common import NewUUID, NodeId
-from shared.types.graphs.common import (
- Edge,
- EdgeData,
-)
-
-
-class EdgeId(NewUUID):
- pass
-
-
-class AddressingProtocol(str, Enum):
- IPvAny = "IPvAny"
-
-
-class ApplicationProtocol(str, Enum):
- MLX = "MLX"
-
-
-AdP = TypeVar("AdP", bound=AddressingProtocol)
-ApP = TypeVar("ApP", bound=ApplicationProtocol)
-
-
-@final
-class EdgeDataTransferRate(BaseModel):
- throughput: float
- latency: float
- jitter: float
-
-
-class NetworkEdgeMetadata(BaseModel, Generic[AdP, ApP]): ...
-
-
-@final
-class NetworkEdgeType(BaseModel, Generic[AdP, ApP]):
- addressing_protocol: AdP
- application_protocol: ApP
-
-
-@final
-class MLXEdgeContext(
- NetworkEdgeMetadata[AddressingProtocol.IPvAny, ApplicationProtocol.MLX]
-):
- source_ip: IPvAnyAddress
- sink_ip: IPvAnyAddress
-
-
-class NetworkEdgeInfoType(str, Enum):
- network_profile = "network_profile"
- other = "other"
-
-
-AllNetworkEdgeInfo = Tuple[NetworkEdgeInfoType.network_profile]
-
-
-NetworkEdgeInfoTypeT = TypeVar(
- "NetworkEdgeInfoTypeT", bound=NetworkEdgeInfoType, covariant=True
-)
-
-
-class NetworkEdgeInfo(BaseModel, Generic[NetworkEdgeInfoTypeT]):
- edge_info_type: NetworkEdgeInfoTypeT
-
-
-SetOfEdgeInfo = TypeVar("SetOfEdgeInfo", bound=Tuple[NetworkEdgeInfoType, ...])
-
-
-class NetworkEdgeData(EdgeData[NetworkEdgeType[AdP, ApP]], Generic[AdP, ApP]):
- edge_info: Mapping[NetworkEdgeInfoType, NetworkEdgeInfo[NetworkEdgeInfoType]]
- edge_metadata: NetworkEdgeMetadata[AdP, ApP]
-
-
-class NetworkEdgeProfile(NetworkEdgeInfo[NetworkEdgeInfoTypeT]):
- edge_data_transfer_rate: EdgeDataTransferRate
-
-
-class NetworkEdge(Edge[NetworkEdgeType[AdP, ApP], EdgeId, NodeId]): ...
diff --git a/shared/types/networking/services.py b/shared/types/networking/services.py
index f7319c43..51620421 100644
--- a/shared/types/networking/services.py
+++ b/shared/types/networking/services.py
@@ -1,10 +1,8 @@
from typing import Callable, NewType, Protocol, TypeVar
-from shared.types.networking.edges import (
- AddressingProtocol,
- ApplicationProtocol,
- EdgeId,
- NetworkEdge,
+from shared.types.networking.control_plane import (
+ ControlPlaneEdgeId,
+ ControlPlaneEdgeType,
)
TopicName = NewType("TopicName", str)
@@ -15,15 +13,12 @@ MessageT = TypeVar("MessageT", bound=object)
PubSubMessageHandler = Callable[[TopicName, MessageT], None]
NodeConnectedHandler = Callable[
[
- EdgeId,
- NetworkEdge[
- AddressingProtocol,
- ApplicationProtocol,
- ],
+ ControlPlaneEdgeId,
+ ControlPlaneEdgeType,
],
None,
]
-NodeDisconnectedHandler = Callable[[EdgeId], None]
+NodeDisconnectedHandler = Callable[[ControlPlaneEdgeId], None]
class DiscoveryService(Protocol):
diff --git a/shared/types/networking/topology.py b/shared/types/networking/topology.py
index 6768c15c..f59d7064 100644
--- a/shared/types/networking/topology.py
+++ b/shared/types/networking/topology.py
@@ -1,40 +1,74 @@
from shared.types.common import NodeId
from shared.types.graphs.common import Graph, GraphData
-from shared.types.networking.edges import (
+from shared.types.networking.control_plane import ControlPlaneEdgeId
+from shared.types.networking.data_plane import (
AddressingProtocol,
ApplicationProtocol,
- EdgeId,
- NetworkEdge,
+ DataPlaneEdge,
+ DataPlaneEdgeId,
)
+from shared.types.worker.common import NodeStatus
-class Topology(
+class DataPlaneTopology(
Graph[
- NetworkEdge[AddressingProtocol, ApplicationProtocol],
+ DataPlaneEdge[AddressingProtocol, ApplicationProtocol],
None,
- EdgeId,
+ DataPlaneEdgeId,
NodeId,
]
):
graph_data: GraphData[
- NetworkEdge[AddressingProtocol, ApplicationProtocol],
+ DataPlaneEdge[AddressingProtocol, ApplicationProtocol],
None,
- EdgeId,
+ DataPlaneEdgeId,
NodeId,
]
-class OrphanedPartOfTopology(
+class OrphanedPartOfDataPlaneTopology(
Graph[
- NetworkEdge[AddressingProtocol, ApplicationProtocol],
+ DataPlaneEdge[AddressingProtocol, ApplicationProtocol],
None,
- EdgeId,
+ DataPlaneEdgeId,
NodeId,
]
):
graph_data: GraphData[
- NetworkEdge[AddressingProtocol, ApplicationProtocol],
+ DataPlaneEdge[AddressingProtocol, ApplicationProtocol],
None,
- EdgeId,
+ DataPlaneEdgeId,
+ NodeId,
+ ]
+
+
+class ControlPlaneTopology(
+ Graph[
+ None,
+ NodeStatus,
+ ControlPlaneEdgeId,
+ NodeId,
+ ]
+):
+ graph_data: GraphData[
+ None,
+ NodeStatus,
+ ControlPlaneEdgeId,
+ NodeId,
+ ]
+
+
+class OrphanedPartOfControlPlaneTopology(
+ Graph[
+ None,
+ NodeStatus,
+ ControlPlaneEdgeId,
+ NodeId,
+ ]
+):
+ graph_data: GraphData[
+ None,
+ NodeStatus,
+ ControlPlaneEdgeId,
NodeId,
]
diff --git a/shared/types/profiling/common.py b/shared/types/profiling/common.py
index 5faffb43..0c09b8f3 100644
--- a/shared/types/profiling/common.py
+++ b/shared/types/profiling/common.py
@@ -1,4 +1,4 @@
from pydantic import BaseModel
-class NodeProfile(BaseModel): ...
+class NodePerformanceProfile(BaseModel): ...
diff --git a/shared/types/states/master.py b/shared/types/states/master.py
index 5f47ec18..b6486a86 100644
--- a/shared/types/states/master.py
+++ b/shared/types/states/master.py
@@ -7,24 +7,28 @@ from pydantic import BaseModel
from shared.types.common import NodeId
from shared.types.events.common import (
- EdgeEventTypes,
+ ControlPlaneEventTypes,
+ DataPlaneEventTypes,
Event,
EventTypes,
- NodeProfileEventTypes,
- NodeStatusEventTypes,
+ NodePerformanceEventTypes,
State,
)
from shared.types.graphs.resource_graph import ResourceGraph
-from shared.types.networking.edges import (
+from shared.types.networking.data_plane import (
AddressingProtocol,
ApplicationProtocol,
- EdgeId,
- NetworkEdge,
+ DataPlaneEdge,
+ DataPlaneEdgeId,
)
-from shared.types.networking.topology import OrphanedPartOfTopology, Topology
-from shared.types.profiling.common import NodeProfile
+from shared.types.networking.topology import (
+ ControlPlaneTopology,
+ DataPlaneTopology,
+ OrphanedPartOfControlPlaneTopology,
+ OrphanedPartOfDataPlaneTopology,
+)
+from shared.types.profiling.common import NodePerformanceProfile
from shared.types.states.shared import SharedState
-from shared.types.worker.common import NodeStatus
from shared.types.worker.instances import InstanceData, InstanceId
@@ -42,28 +46,33 @@ class CachePolicy(BaseModel, Generic[CachePolicyTypeT]):
policy_type: CachePolicyTypeT
-class NodeProfileState(State[NodeProfileEventTypes]):
- node_profiles: Mapping[NodeId, NodeProfile]
+class NodePerformanceProfileState(State[NodePerformanceEventTypes]):
+ node_profiles: Mapping[NodeId, NodePerformanceProfile]
+
+class DataPlaneNetworkState(State[DataPlaneEventTypes]):
+ topology: DataPlaneTopology
+ history: Sequence[OrphanedPartOfDataPlaneTopology]
-class NodeStatusState(State[NodeStatusEventTypes]):
- node_status: Mapping[NodeId, NodeStatus]
+ def delete_edge(self, edge_id: DataPlaneEdgeId) -> None: ...
+ def add_edge(
+ self, edge: DataPlaneEdge[AddressingProtocol, ApplicationProtocol]
+ ) -> None: ...
-class NetworkState(State[EdgeEventTypes]):
- topology: Topology
- history: Sequence[OrphanedPartOfTopology]
+class ControlPlaneNetworkState(State[ControlPlaneEventTypes]):
+ topology: ControlPlaneTopology
+ history: Sequence[OrphanedPartOfControlPlaneTopology]
- def delete_edge(self, edge_id: EdgeId) -> None: ...
+ def delete_edge(self, edge_id: DataPlaneEdgeId) -> None: ...
def add_edge(
- self, edge: NetworkEdge[AddressingProtocol, ApplicationProtocol]
+ self, edge: DataPlaneEdge[AddressingProtocol, ApplicationProtocol]
) -> None: ...
class MasterState(SharedState):
- network_state: NetworkState
- node_profiles: NodeProfileState
- node_status: NodeStatusState
+ data_plane_network_state: DataPlaneNetworkState
+ control_plane_network_state: ControlPlaneNetworkState
job_inbox: Queue[ExternalCommand]
job_outbox: Queue[ExternalCommand]
cache_policy: CachePolicy[CachePolicyType]
diff --git a/shared/types/states/shared.py b/shared/types/states/shared.py
index 1dae6823..c06c8c21 100644
--- a/shared/types/states/shared.py
+++ b/shared/types/states/shared.py
@@ -10,7 +10,7 @@ from shared.types.worker.common import InstanceId
from shared.types.worker.instances import BaseInstance
-class Instances(State[InstanceStateEventTypes]):
+class KnownInstances(State[InstanceStateEventTypes]):
instances: Mapping[InstanceId, BaseInstance]
@@ -20,5 +20,11 @@ class Tasks(State[TaskEventTypes]):
class SharedState(BaseModel):
node_id: NodeId
- known_instances: Instances
+ known_instances: KnownInstances
compute_tasks: Sequence[Task[TaskType]]
+
+ def get_node_id(self) -> NodeId: ...
+
+ def get_tasks_by_instance(
+ self, instance_id: InstanceId
+ ) -> Sequence[Task[TaskType]]: ...
diff --git a/shared/types/states/worker.py b/shared/types/states/worker.py
index 5db788df..02b1fb67 100644
--- a/shared/types/states/worker.py
+++ b/shared/types/states/worker.py
@@ -2,14 +2,14 @@ from collections.abc import Mapping
from shared.types.common import NodeId
from shared.types.events.common import (
- NodeStatusEventTypes,
+ ControlPlaneEventTypes,
State,
)
from shared.types.states.shared import SharedState
from shared.types.worker.common import NodeStatus
-class NodeStatusState(State[NodeStatusEventTypes]):
+class NodeStatusState(State[ControlPlaneEventTypes]):
node_status: Mapping[NodeId, NodeStatus]
diff --git a/shared/types/tasks/common.py b/shared/types/tasks/common.py
index da99804f..a01c641d 100644
--- a/shared/types/tasks/common.py
+++ b/shared/types/tasks/common.py
@@ -1,6 +1,6 @@
from collections.abc import Mapping
from enum import Enum
-from typing import Any, Generic, Literal, TypeVar, Union
+from typing import Generic, Literal, TypeVar, Union
import openai.types.chat as openai
from pydantic import BaseModel
@@ -23,7 +23,6 @@ TaskTypeT = TypeVar("TaskTypeT", bound=TaskType)
class TaskData(BaseModel, Generic[TaskTypeT]):
task_type: TaskTypeT
- task_data: Any
class ChatCompletionNonStreamingTask(TaskData[TaskType.ChatCompletionNonStreaming]):
@@ -52,33 +51,36 @@ TaskStatusTypeT = TypeVar(
)
-class TaskUpdate(BaseModel, Generic[TaskStatusTypeT]):
+class TaskArtifact(BaseModel, Generic[TaskTypeT]): ...
+
+
+class TaskUpdate(BaseModel, Generic[TaskStatusTypeT, TaskTypeT]):
task_status: TaskStatusTypeT
-class PendingTask(TaskUpdate[TaskStatusType.Pending]):
+class PendingTask(TaskUpdate[TaskStatusType.Pending, TaskTypeT]):
task_status: Literal[TaskStatusType.Pending]
-class RunningTask(TaskUpdate[TaskStatusType.Running]):
+class RunningTask(TaskUpdate[TaskStatusType.Running, TaskTypeT]):
task_status: Literal[TaskStatusType.Running]
-class CompletedTask(TaskUpdate[TaskStatusType.Complete]):
+class CompletedTask(TaskUpdate[TaskStatusType.Complete, TaskTypeT]):
task_status: Literal[TaskStatusType.Complete]
- task_artifact: bytes
+ task_artifact: TaskArtifact[TaskTypeT]
-class FailedTask(TaskUpdate[TaskStatusType.Failed]):
+class FailedTask(TaskUpdate[TaskStatusType.Failed, TaskTypeT]):
task_status: Literal[TaskStatusType.Failed]
error_message: Mapping[RunnerId, str]
-class BaseTask(BaseModel):
- task_data: TaskData[TaskType]
- task_status: TaskUpdate[TaskStatusType]
+class BaseTask(BaseModel, Generic[TaskTypeT]):
+ task_data: TaskData[TaskTypeT]
+ task_status: TaskUpdate[TaskStatusType, TaskTypeT]
on_instance: InstanceId
-class Task(BaseTask):
+class Task(BaseTask[TaskTypeT]):
task_id: TaskId
← c4569343 refactor: Remove timestamp from Wrapped Events
·
back to Exo
·
fix: Use state for tasks cda3de2a →