← back to Exo
Topology update
cd9a1a9192580b458e73f1dec40c382d541586ba · 2025-07-22 22:29:17 +0100 · Seth Howes
Files touched
M engines/mlx/utils_mlx.pyM master/api.pyM master/main.pyM master/placement.pyA master/tests/test_topology.pyD shared/graphs.pyM shared/tests/test_sqlite_connector.pyA shared/topology.pyM shared/types/events/chunks.pyM shared/types/events/events.pyD shared/types/graphs/common.pyD shared/types/graphs/topology.pyA shared/types/profiling.pyD shared/types/profiling/common.pyR100 shared/types/tasks/request.py shared/types/request.pyM shared/types/state.pyR085 shared/types/tasks/common.py shared/types/tasks.pyA shared/types/topology.pyM shared/types/worker/commands_runner.pyM shared/types/worker/ops.pyM shared/types/worker/resource_monitor.pyM worker/runner/runner.pyM worker/runner/runner_supervisor.pyM worker/tests/conftest.pyM worker/tests/test_serdes.pyM worker/tests/test_supervisor.pyM worker/tests/test_worker_handlers.py
Diff
commit cd9a1a9192580b458e73f1dec40c382d541586ba
Author: Seth Howes <71157822+sethhowes@users.noreply.github.com>
Date: Tue Jul 22 22:29:17 2025 +0100
Topology update
---
engines/mlx/utils_mlx.py | 2 +-
master/api.py | 4 +-
master/main.py | 2 +-
master/placement.py | 7 +-
master/tests/test_topology.py | 170 +++++++++++++++++++
shared/graphs.py | 261 -----------------------------
shared/tests/test_sqlite_connector.py | 2 +-
shared/topology.py | 109 ++++++++++++
shared/types/events/chunks.py | 2 +-
shared/types/events/events.py | 22 +--
shared/types/graphs/common.py | 172 -------------------
shared/types/graphs/topology.py | 48 ------
shared/types/profiling.py | 32 ++++
shared/types/profiling/common.py | 54 ------
shared/types/{tasks => }/request.py | 0
shared/types/state.py | 26 +--
shared/types/{tasks/common.py => tasks.py} | 5 +-
shared/types/topology.py | 68 ++++++++
shared/types/worker/commands_runner.py | 2 +-
shared/types/worker/ops.py | 2 +-
shared/types/worker/resource_monitor.py | 29 +---
worker/runner/runner.py | 2 +-
worker/runner/runner_supervisor.py | 7 +-
worker/tests/conftest.py | 6 +-
worker/tests/test_serdes.py | 2 +-
worker/tests/test_supervisor.py | 4 +-
worker/tests/test_worker_handlers.py | 2 +-
27 files changed, 426 insertions(+), 616 deletions(-)
diff --git a/engines/mlx/utils_mlx.py b/engines/mlx/utils_mlx.py
index 52777c53..781c76f9 100644
--- a/engines/mlx/utils_mlx.py
+++ b/engines/mlx/utils_mlx.py
@@ -12,7 +12,7 @@ from mlx_lm.utils import load_model # type: ignore
from pydantic import RootModel
from engines.mlx.auto_parallel import auto_parallel
-from shared.types.tasks.common import ChatCompletionTaskParams
+from shared.types.tasks import ChatCompletionTaskParams
from shared.types.worker.mlx import Host
from shared.types.worker.shards import ShardMetadata
from worker.download.download_utils import build_model_path
diff --git a/master/api.py b/master/api.py
index 219b5f57..28c78e48 100644
--- a/master/api.py
+++ b/master/api.py
@@ -14,8 +14,8 @@ from shared.types.events.chunks import TokenChunk
from shared.types.events.components import EventFromEventLog
from shared.types.events.events import ChunkGenerated
from shared.types.events.registry import Event
-from shared.types.tasks.common import ChatCompletionTaskParams
-from shared.types.tasks.request import APIRequest, RequestId
+from shared.types.request import APIRequest, RequestId
+from shared.types.tasks import ChatCompletionTaskParams
class Message(BaseModel):
diff --git a/master/main.py b/master/main.py
index 37949c27..cb59ec45 100644
--- a/master/main.py
+++ b/master/main.py
@@ -10,7 +10,7 @@ from shared.db.sqlite.event_log_manager import EventLogManager
from shared.types.common import NodeId
from shared.types.events.chunks import TokenChunk
from shared.types.events.events import ChunkGenerated
-from shared.types.tasks.request import APIRequest, RequestId
+from shared.types.request import APIRequest, RequestId
## TODO: Hook this up properly
diff --git a/master/placement.py b/master/placement.py
index 9803816f..be0d8f41 100644
--- a/master/placement.py
+++ b/master/placement.py
@@ -1,14 +1,14 @@
from queue import Queue
from typing import Mapping, Sequence
+from shared.topology import Topology
from shared.types.events.registry import Event
-from shared.types.graphs.topology import Topology
from shared.types.state import CachePolicy
-from shared.types.tasks.common import Task
+from shared.types.tasks import Task
from shared.types.worker.instances import InstanceId, InstanceParams
-def get_instance_placement(
+def get_instance_placements(
inbox: Queue[Task],
outbox: Queue[Task],
topology: Topology,
@@ -17,6 +17,7 @@ def get_instance_placement(
) -> Mapping[InstanceId, InstanceParams]: ...
+
def get_transition_events(
current_instances: Mapping[InstanceId, InstanceParams],
target_instances: Mapping[InstanceId, InstanceParams],
diff --git a/master/tests/test_topology.py b/master/tests/test_topology.py
new file mode 100644
index 00000000..5eaca934
--- /dev/null
+++ b/master/tests/test_topology.py
@@ -0,0 +1,170 @@
+import pytest
+
+from shared.topology import Topology
+from shared.types.profiling import (
+ MemoryPerformanceProfile,
+ NodePerformanceProfile,
+ SystemPerformanceProfile,
+)
+from shared.types.topology import Connection, ConnectionProfile, Node, NodeId
+
+
+@pytest.fixture
+def topology() -> Topology:
+ return Topology()
+
+@pytest.fixture
+def connection() -> Connection:
+ return Connection(source_node_id=NodeId(), sink_node_id=NodeId(), source_multiaddr="/ip4/127.0.0.1/tcp/1234", sink_multiaddr="/ip4/127.0.0.1/tcp/1235", connection_profile=ConnectionProfile(throughput=1000, latency=1000, jitter=1000))
+
+@pytest.fixture
+def node_profile() -> NodePerformanceProfile:
+ memory_profile = MemoryPerformanceProfile(ram_total=1000, ram_used=0, swap_total=1000, swap_used=0)
+ system_profile = SystemPerformanceProfile(flops_fp16=1000)
+ return NodePerformanceProfile(model_id="test", chip_id="test", memory=memory_profile, network_interfaces=[], system=system_profile)
+
+@pytest.fixture
+def connection_profile() -> ConnectionProfile:
+ return ConnectionProfile(throughput=1000, latency=1000, jitter=1000)
+
+def test_add_node(topology: Topology, node_profile: NodePerformanceProfile):
+ # arrange
+ node_id = NodeId()
+
+ # act
+ topology.add_node(Node(node_id=node_id, node_profile=node_profile), node_id=node_id)
+
+ # assert
+ data = topology.get_node_profile(node_id)
+ assert data == node_profile
+
+
+def test_add_connection(topology: Topology, node_profile: NodePerformanceProfile, connection: Connection):
+ # arrange
+ topology.add_node(Node(node_id=connection.source_node_id, node_profile=node_profile), node_id=connection.source_node_id)
+ topology.add_node(Node(node_id=connection.sink_node_id, node_profile=node_profile), node_id=connection.sink_node_id)
+ topology.add_connection(connection)
+
+ # act
+ data = topology.get_connection_profile(connection)
+
+ # assert
+ assert data == connection.connection_profile
+
+def test_update_node_profile(topology: Topology, node_profile: NodePerformanceProfile, connection: Connection):
+ # arrange
+ topology.add_node(Node(node_id=connection.source_node_id, node_profile=node_profile), node_id=connection.source_node_id)
+ topology.add_node(Node(node_id=connection.sink_node_id, node_profile=node_profile), node_id=connection.sink_node_id)
+ topology.add_connection(connection)
+
+ new_node_profile = NodePerformanceProfile(model_id="test", chip_id="test", memory=MemoryPerformanceProfile(ram_total=1000, ram_used=0, swap_total=1000, swap_used=0), network_interfaces=[], system=SystemPerformanceProfile(flops_fp16=1000))
+
+ # act
+ topology.update_node_profile(connection.source_node_id, node_profile=new_node_profile)
+
+ # assert
+ data = topology.get_node_profile(connection.source_node_id)
+ assert data == new_node_profile
+
+def test_update_connection_profile(topology: Topology, node_profile: NodePerformanceProfile, connection: Connection):
+ # arrange
+ topology.add_node(Node(node_id=connection.source_node_id, node_profile=node_profile), node_id=connection.source_node_id)
+ topology.add_node(Node(node_id=connection.sink_node_id, node_profile=node_profile), node_id=connection.sink_node_id)
+ topology.add_connection(connection)
+
+ new_connection_profile = ConnectionProfile(throughput=2000, latency=2000, jitter=2000)
+ connection = Connection(source_node_id=connection.source_node_id, sink_node_id=connection.sink_node_id, source_multiaddr=connection.source_multiaddr, sink_multiaddr=connection.sink_multiaddr, connection_profile=new_connection_profile)
+
+ # act
+ topology.update_connection_profile(connection)
+
+ # assert
+ data = topology.get_connection_profile(connection)
+ assert data == new_connection_profile
+
+def test_remove_connection_still_connected(topology: Topology, node_profile: NodePerformanceProfile, connection: Connection):
+ # arrange
+ topology.add_node(Node(node_id=connection.source_node_id, node_profile=node_profile), node_id=connection.source_node_id)
+ topology.add_node(Node(node_id=connection.sink_node_id, node_profile=node_profile), node_id=connection.sink_node_id)
+ topology.add_connection(connection)
+
+ # act
+ topology.remove_connection(connection)
+
+ # assert
+ with pytest.raises(IndexError):
+ topology.get_connection_profile(connection)
+
+def test_remove_connection_bridge(topology: Topology, node_profile: NodePerformanceProfile, connection: Connection):
+ """Create a bridge scenario: master -> node_a -> node_b
+ and remove the bridge connection (master -> node_a)"""
+ # arrange
+ master_id = NodeId()
+ node_a_id = NodeId()
+ node_b_id = NodeId()
+
+ topology.add_node(Node(node_id=master_id, node_profile=node_profile), node_id=master_id)
+ topology.add_node(Node(node_id=node_a_id, node_profile=node_profile), node_id=node_a_id)
+ topology.add_node(Node(node_id=node_b_id, node_profile=node_profile), node_id=node_b_id)
+
+ connection_master_to_a = Connection(
+ source_node_id=master_id,
+ sink_node_id=node_a_id,
+ source_multiaddr="/ip4/127.0.0.1/tcp/1234",
+ sink_multiaddr="/ip4/127.0.0.1/tcp/1235",
+ connection_profile=ConnectionProfile(throughput=1000, latency=1000, jitter=1000)
+ )
+
+ connection_a_to_b = Connection(
+ source_node_id=node_a_id,
+ sink_node_id=node_b_id,
+ source_multiaddr="/ip4/127.0.0.1/tcp/1236",
+ sink_multiaddr="/ip4/127.0.0.1/tcp/1237",
+ connection_profile=ConnectionProfile(throughput=1000, latency=1000, jitter=1000)
+ )
+
+ topology.add_connection(connection_master_to_a)
+ topology.add_connection(connection_a_to_b)
+
+ assert len(list(topology.list_nodes())) == 3
+
+ topology.remove_connection(connection_master_to_a)
+
+ remaining_nodes = list(topology.list_nodes())
+ assert len(remaining_nodes) == 1
+ assert remaining_nodes[0].node_id == master_id
+
+ with pytest.raises(KeyError):
+ topology.get_node_profile(node_a_id)
+
+ with pytest.raises(KeyError):
+ topology.get_node_profile(node_b_id)
+
+
+def test_remove_node_still_connected(topology: Topology, node_profile: NodePerformanceProfile, connection: Connection):
+ # arrange
+ topology.add_node(Node(node_id=connection.source_node_id, node_profile=node_profile), node_id=connection.source_node_id)
+ topology.add_node(Node(node_id=connection.sink_node_id, node_profile=node_profile), node_id=connection.sink_node_id)
+ topology.add_connection(connection)
+
+ # act
+ topology.remove_node(connection.source_node_id)
+
+ # assert
+ with pytest.raises(KeyError):
+ topology.get_node_profile(connection.source_node_id)
+
+
+def test_list_nodes(topology: Topology, node_profile: NodePerformanceProfile, connection: Connection):
+ # arrange
+ topology.add_node(Node(node_id=connection.source_node_id, node_profile=node_profile), node_id=connection.source_node_id)
+ topology.add_node(Node(node_id=connection.sink_node_id, node_profile=node_profile), node_id=connection.sink_node_id)
+ topology.add_connection(connection)
+
+ # act
+ nodes = list(topology.list_nodes())
+
+ # assert
+ assert len(nodes) == 2
+ assert all(isinstance(node, Node) for node in nodes)
+ assert {node.node_id for node in nodes} == {connection.source_node_id, connection.sink_node_id}
diff --git a/shared/graphs.py b/shared/graphs.py
deleted file mode 100644
index a48474da..00000000
--- a/shared/graphs.py
+++ /dev/null
@@ -1,261 +0,0 @@
-from dataclasses import dataclass
-from typing import Any, Callable, Mapping, Set
-
-import rustworkx as rx
-from pydantic import TypeAdapter
-from pydantic_core import core_schema
-
-from shared.types.graphs.common import (
- Edge,
- EdgeData,
- EdgeIdT,
- EdgeTypeT,
- MutableGraphProtocol,
- Vertex,
- VertexData,
- VertexIdT,
- VertexTypeT,
-)
-from shared.types.graphs.pydantic import PydanticGraph
-
-
-@dataclass(frozen=True)
-class _VertexWrapper[VertexTypeT, VertexIdT]:
- """Internal wrapper to store vertex ID alongside vertex data."""
-
- vertex_id: VertexIdT
- vertex_data: VertexData[VertexTypeT]
-
-
-@dataclass(frozen=True)
-class _EdgeWrapper[EdgeTypeT, EdgeIdT]:
- """Internal wrapper to store edge ID alongside edge data."""
-
- edge_id: EdgeIdT
- edge_data: EdgeData[EdgeTypeT]
-
-
-class Graph(MutableGraphProtocol[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]):
- edge_base: TypeAdapter[EdgeTypeT]
- vertex_base: TypeAdapter[VertexTypeT]
-
- _graph: rx.PyDiGraph[
- _VertexWrapper[VertexTypeT, VertexIdT], _EdgeWrapper[EdgeTypeT, EdgeIdT]
- ]
- _vertex_id_to_index: dict[VertexIdT, int]
- _edge_id_to_endpoints: dict[EdgeIdT, tuple[int, int]]
-
- def __init__(
- self, edge_base: TypeAdapter[EdgeTypeT], vertex_base: TypeAdapter[VertexTypeT]
- ) -> None:
- self.edge_base = edge_base
- self.vertex_base = vertex_base
- self._graph = rx.PyDiGraph()
- self._vertex_id_to_index = {}
- self._edge_id_to_endpoints = {}
-
- # TODO: I'm not sure if this is the right thing, but we'll simplify the graph stuff anyway so fine for now.
- @classmethod
- def __get_pydantic_core_schema__(
- cls,
- _source: type[Any],
- handler: Callable[[Any], core_schema.CoreSchema],
- ) -> core_schema.CoreSchema:
- pydantic_graph_schema = handler(PydanticGraph)
-
- def serializer(
- instance: "Graph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]",
- ) -> dict[str, Any]:
- return {
- "vertices": list(instance.get_vertex_data(instance.list_vertices())),
- "edges": list(instance.get_edge_data(instance.list_edges())),
- }
-
- return core_schema.json_or_python_schema(
- json_schema=pydantic_graph_schema,
- python_schema=core_schema.no_info_plain_validator_function(cls.validate),
- serialization=core_schema.plain_serializer_function_ser_schema(serializer),
- )
-
- @classmethod
- def validate(
- cls, value: Any # type: ignore
- ) -> "Graph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]":
- if isinstance(value, cls):
- return value
-
- if isinstance(value, dict):
- raise NotImplementedError(
- "Deserializing a Graph from a dictionary is not yet supported. "
- "Please initialize the Graph object directly."
- )
-
- raise TypeError("Unsupported type for Graph validation")
-
- ###
- # GraphProtocol methods
- ###
-
- def list_edges(self) -> Set[EdgeIdT]:
- return {edge.edge_id for edge in self._graph.edges()}
-
- def list_vertices(self) -> Set[VertexIdT]:
- return {node.vertex_id for node in self._graph.nodes()}
-
- def get_vertices_from_edges(
- self, edges: Set[EdgeIdT]
- ) -> Mapping[EdgeIdT, Set[VertexIdT]]:
- result: dict[EdgeIdT, Set[VertexIdT]] = {}
-
- for edge_id in edges:
- if edge_id in self._edge_id_to_endpoints:
- u_idx, v_idx = self._edge_id_to_endpoints[edge_id]
- u_data = self._graph.get_node_data(u_idx)
- v_data = self._graph.get_node_data(v_idx)
- result[edge_id] = {u_data.vertex_id, v_data.vertex_id}
-
- return result
-
- def get_edges_from_vertices(
- self, vertices: Set[VertexIdT]
- ) -> Mapping[VertexIdT, Set[EdgeIdT]]:
- result: dict[VertexIdT, Set[EdgeIdT]] = {}
-
- for vertex_id in vertices:
- if vertex_id in self._vertex_id_to_index:
- vertex_idx = self._vertex_id_to_index[vertex_id]
- edge_ids: Set[EdgeIdT] = set()
-
- # Get outgoing edges
- for _, _, edge_data in self._graph.out_edges(vertex_idx):
- edge_ids.add(edge_data.edge_id)
-
- # Get incoming edges
- for _, _, edge_data in self._graph.in_edges(vertex_idx):
- edge_ids.add(edge_data.edge_id)
-
- result[vertex_id] = edge_ids
-
- return result
-
- def get_edge_data(
- self, edges: Set[EdgeIdT]
- ) -> Mapping[EdgeIdT, EdgeData[EdgeTypeT]]:
- result: dict[EdgeIdT, EdgeData[EdgeTypeT]] = {}
-
- for edge_id in edges:
- if edge_id in self._edge_id_to_endpoints:
- u_idx, v_idx = self._edge_id_to_endpoints[edge_id]
- edge_wrapper = self._graph.get_edge_data(u_idx, v_idx)
- result[edge_id] = edge_wrapper.edge_data
-
- return result
-
- def get_vertex_data(
- self, vertices: Set[VertexIdT]
- ) -> Mapping[VertexIdT, VertexData[VertexTypeT]]:
- result: dict[VertexIdT, VertexData[VertexTypeT]] = {}
-
- for vertex_id in vertices:
- if vertex_id in self._vertex_id_to_index:
- vertex_idx = self._vertex_id_to_index[vertex_id]
- vertex_wrapper = self._graph.get_node_data(vertex_idx)
- result[vertex_id] = vertex_wrapper.vertex_data
-
- return result
-
- ###
- # MutableGraphProtocol methods
- ###
-
- def check_edges_exists(self, edge_id: EdgeIdT) -> bool:
- return edge_id in self._edge_id_to_endpoints
-
- def check_vertex_exists(self, vertex_id: VertexIdT) -> bool:
- return vertex_id in self._vertex_id_to_index
-
- def _add_edge(self, edge_id: EdgeIdT, edge_data: EdgeData[EdgeTypeT]) -> None:
- # This internal method is not used in favor of a safer `attach_edge` implementation.
- raise NotImplementedError(
- "Use attach_edge to add edges. The internal _add_edge protocol method is flawed."
- )
-
- def _add_vertex(
- self, vertex_id: VertexIdT, vertex_data: VertexData[VertexTypeT]
- ) -> None:
- if vertex_id not in self._vertex_id_to_index:
- wrapper = _VertexWrapper(vertex_id=vertex_id, vertex_data=vertex_data)
- idx = self._graph.add_node(wrapper)
- self._vertex_id_to_index[vertex_id] = idx
-
- def _remove_edge(self, edge_id: EdgeIdT) -> None:
- if edge_id in self._edge_id_to_endpoints:
- u_idx, v_idx = self._edge_id_to_endpoints[edge_id]
- self._graph.remove_edge(u_idx, v_idx)
- del self._edge_id_to_endpoints[edge_id]
- else:
- raise ValueError(f"Edge with id {edge_id} not found.")
-
- def _remove_vertex(self, vertex_id: VertexIdT) -> None:
- if vertex_id in self._vertex_id_to_index:
- vertex_idx = self._vertex_id_to_index[vertex_id]
-
- # Remove any edges connected to this vertex from our mapping
- edges_to_remove: list[EdgeIdT] = []
- for edge_id, (u_idx, v_idx) in self._edge_id_to_endpoints.items():
- if u_idx == vertex_idx or v_idx == vertex_idx:
- edges_to_remove.append(edge_id)
-
- for edge_id in edges_to_remove:
- del self._edge_id_to_endpoints[edge_id]
-
- # Remove the vertex from the graph
- self._graph.remove_node(vertex_idx)
- del self._vertex_id_to_index[vertex_id]
- else:
- raise ValueError(f"Vertex with id {vertex_id} not found.")
-
- def attach_edge(
- self,
- edge: Edge[EdgeTypeT, EdgeIdT, VertexIdT],
- extra_vertex: Vertex[VertexTypeT, EdgeIdT, VertexIdT] | None = None,
- ) -> None:
- """
- Attaches an edge to the graph, overriding the default protocol implementation.
-
- This implementation corrects a flaw in the protocol's `_add_edge`
- signature and provides more intuitive behavior when connecting existing vertices.
- """
- base_vertex_id, target_vertex_id = edge.edge_vertices
-
- if not self.check_vertex_exists(base_vertex_id):
- raise ValueError(f"Base vertex {base_vertex_id} does not exist.")
-
- target_vertex_exists = self.check_vertex_exists(target_vertex_id)
-
- if not target_vertex_exists:
- if extra_vertex is None:
- raise ValueError(
- f"Target vertex {target_vertex_id} does not exist and no `extra_vertex` was provided."
- )
- if extra_vertex.vertex_id != target_vertex_id:
- raise ValueError(
- f"The ID of `extra_vertex` ({extra_vertex.vertex_id}) does not match "
- f"the target vertex ID of the edge ({target_vertex_id})."
- )
- self._add_vertex(extra_vertex.vertex_id, extra_vertex.vertex_data)
- elif extra_vertex is not None:
- raise ValueError(
- f"Target vertex {target_vertex_id} already exists, but `extra_vertex` was provided."
- )
-
- # Get the internal indices
- base_idx = self._vertex_id_to_index[base_vertex_id]
- target_idx = self._vertex_id_to_index[target_vertex_id]
-
- # Create edge wrapper and add to graph
- edge_wrapper = _EdgeWrapper(edge_id=edge.edge_id, edge_data=edge.edge_data)
- self._graph.add_edge(base_idx, target_idx, edge_wrapper)
-
- # Store the mapping
- self._edge_id_to_endpoints[edge.edge_id] = (base_idx, target_idx)
diff --git a/shared/tests/test_sqlite_connector.py b/shared/tests/test_sqlite_connector.py
index 7bd98b40..32e9ea8c 100644
--- a/shared/tests/test_sqlite_connector.py
+++ b/shared/tests/test_sqlite_connector.py
@@ -16,7 +16,7 @@ from shared.types.events.events import (
ChunkGenerated,
EventType,
)
-from shared.types.tasks.request import RequestId
+from shared.types.request import RequestId
# Type ignore comment for all protected member access in this test file
# pyright: reportPrivateUsage=false
diff --git a/shared/topology.py b/shared/topology.py
new file mode 100644
index 00000000..289912f3
--- /dev/null
+++ b/shared/topology.py
@@ -0,0 +1,109 @@
+from typing import Iterable
+
+import rustworkx as rx
+
+from shared.types.common import NodeId
+from shared.types.profiling import ConnectionProfile, NodePerformanceProfile
+from shared.types.topology import Connection, Node, TopologyProto
+
+
+class Topology(TopologyProto):
+ def __init__(self) -> None:
+ self._graph: rx.PyDiGraph[Node, Connection] = rx.PyDiGraph()
+ self._node_id_to_rx_id_map: dict[NodeId, int] = dict()
+ self._rx_id_to_node_id_map: dict[int, NodeId] = dict()
+ self._edge_id_to_rx_id_map: dict[Connection, int] = dict()
+ self.master_node_id: NodeId | None = None
+
+ # TODO: implement serialization + deserialization method
+
+ def add_node(self, node: Node, node_id: NodeId) -> None:
+ if node_id in self._node_id_to_rx_id_map:
+ raise ValueError("Node already exists")
+ rx_id = self._graph.add_node(node)
+ self._node_id_to_rx_id_map[node_id] = rx_id
+ self._rx_id_to_node_id_map[rx_id] = node_id
+
+
+ def add_connection(
+ self,
+ connection: Connection,
+ ) -> None:
+ if connection.source_node_id not in self._node_id_to_rx_id_map:
+ self.add_node(Node(node_id=connection.source_node_id), node_id=connection.source_node_id)
+ if connection.sink_node_id not in self._node_id_to_rx_id_map:
+ self.add_node(Node(node_id=connection.sink_node_id), node_id=connection.sink_node_id)
+
+ src_id = self._node_id_to_rx_id_map[connection.source_node_id]
+ sink_id = self._node_id_to_rx_id_map[connection.sink_node_id]
+
+ rx_id = self._graph.add_edge(src_id, sink_id, connection)
+ self._edge_id_to_rx_id_map[connection] = rx_id
+
+ def list_nodes(self) -> Iterable[Node]:
+ yield from (self._graph[i] for i in self._graph.node_indices())
+
+ def list_connections(self) -> Iterable[Connection]:
+ for (_, _, connection) in self._graph.weighted_edge_list():
+ yield connection
+
+ def get_node_profile(self, node_id: NodeId) -> NodePerformanceProfile | None:
+ rx_idx = self._node_id_to_rx_id_map[node_id]
+ return self._graph.get_node_data(rx_idx).node_profile
+
+ def update_node_profile(self, node_id: NodeId, node_profile: NodePerformanceProfile) -> None:
+ rx_idx = self._node_id_to_rx_id_map[node_id]
+ self._graph[rx_idx].node_profile = node_profile
+
+ def update_connection_profile(self, connection: Connection) -> None:
+ rx_idx = self._edge_id_to_rx_id_map[connection]
+ self._graph.update_edge_by_index(rx_idx, connection)
+
+ def get_connection_profile(self, connection: Connection) -> ConnectionProfile | None:
+ rx_idx = self._edge_id_to_rx_id_map[connection]
+ return self._graph.get_edge_data_by_index(rx_idx).connection_profile
+
+ def remove_node(self, node_id: NodeId) -> None:
+ rx_idx = self._node_id_to_rx_id_map[node_id]
+ self._graph.remove_node(rx_idx)
+
+ del self._node_id_to_rx_id_map[node_id]
+ del self._rx_id_to_node_id_map[rx_idx]
+
+ def remove_connection(self, connection: Connection) -> None:
+ rx_idx = self._edge_id_to_rx_id_map[connection]
+ if self._is_bridge(connection):
+ orphan_node_ids = self._get_orphan_node_ids(connection.source_node_id, connection)
+ for orphan_node_id in orphan_node_ids:
+ orphan_node_rx_id = self._node_id_to_rx_id_map[orphan_node_id]
+ self._graph.remove_node(orphan_node_rx_id)
+ del self._node_id_to_rx_id_map[orphan_node_id]
+ del self._rx_id_to_node_id_map[orphan_node_rx_id]
+ else:
+ self._graph.remove_edge_from_index(rx_idx)
+ del self._edge_id_to_rx_id_map[connection]
+ del self._rx_id_to_node_id_map[rx_idx]
+
+ def _is_bridge(self, connection: Connection) -> bool:
+ edge_idx = self._edge_id_to_rx_id_map[connection]
+ graph_copy = self._graph.copy().to_undirected()
+ components_before = rx.number_connected_components(graph_copy)
+
+ graph_copy.remove_edge_from_index(edge_idx)
+ components_after = rx.number_connected_components(graph_copy)
+
+ return components_after > components_before
+
+ def _get_orphan_node_ids(self, master_node_id: NodeId, connection: Connection) -> list[NodeId]:
+ edge_idx = self._edge_id_to_rx_id_map[connection]
+ graph_copy = self._graph.copy().to_undirected()
+ graph_copy.remove_edge_from_index(edge_idx)
+ components = rx.connected_components(graph_copy)
+
+ orphan_node_rx_ids: set[int] = set()
+ master_node_rx_id = self._node_id_to_rx_id_map[master_node_id]
+ for component in components:
+ if master_node_rx_id not in component:
+ orphan_node_rx_ids.update(component)
+
+ return [self._rx_id_to_node_id_map[rx_id] for rx_id in orphan_node_rx_ids]
diff --git a/shared/types/events/chunks.py b/shared/types/events/chunks.py
index 860633e1..9504496a 100644
--- a/shared/types/events/chunks.py
+++ b/shared/types/events/chunks.py
@@ -5,7 +5,7 @@ from pydantic import BaseModel, Field, TypeAdapter
from shared.openai_compat import FinishReason
from shared.types.models import ModelId
-from shared.types.tasks.request import RequestId
+from shared.types.request import RequestId
class ChunkType(str, Enum):
diff --git a/shared/types/events/events.py b/shared/types/events/events.py
index dd9a1d5c..90c98a27 100644
--- a/shared/types/events/events.py
+++ b/shared/types/events/events.py
@@ -2,6 +2,7 @@ from __future__ import annotations
from typing import Literal, Tuple
+from shared.topology import Connection, ConnectionProfile, Node, NodePerformanceProfile
from shared.types.common import NodeId
from shared.types.events.chunks import GenerationChunk
from shared.types.events.common import (
@@ -9,15 +10,8 @@ from shared.types.events.common import (
EventType,
TimerId,
)
-from shared.types.graphs.topology import (
- TopologyEdge,
- TopologyEdgeId,
- TopologyEdgeProfile,
- TopologyNode,
-)
-from shared.types.profiling.common import NodePerformanceProfile
-from shared.types.tasks.common import Task, TaskId, TaskStatus
-from shared.types.tasks.request import RequestId
+from shared.types.request import RequestId
+from shared.types.tasks import Task, TaskId, TaskStatus
from shared.types.worker.common import InstanceId, NodeStatus
from shared.types.worker.instances import InstanceParams, TypeOfInstance
from shared.types.worker.runners import RunnerId, RunnerStatus
@@ -96,7 +90,7 @@ class NodePerformanceMeasured(BaseEvent[EventType.NodePerformanceMeasured]):
class WorkerConnected(BaseEvent[EventType.WorkerConnected]):
event_type: Literal[EventType.WorkerConnected] = EventType.WorkerConnected
- edge: TopologyEdge
+ edge: Connection
class WorkerStatusUpdated(BaseEvent[EventType.WorkerStatusUpdated]):
@@ -118,18 +112,18 @@ class ChunkGenerated(BaseEvent[EventType.ChunkGenerated]):
class TopologyEdgeCreated(BaseEvent[EventType.TopologyEdgeCreated]):
event_type: Literal[EventType.TopologyEdgeCreated] = EventType.TopologyEdgeCreated
- vertex: TopologyNode
+ vertex: Node
class TopologyEdgeReplacedAtomically(BaseEvent[EventType.TopologyEdgeReplacedAtomically]):
event_type: Literal[EventType.TopologyEdgeReplacedAtomically] = EventType.TopologyEdgeReplacedAtomically
- edge_id: TopologyEdgeId
- edge_profile: TopologyEdgeProfile
+ edge: Connection
+ edge_profile: ConnectionProfile
class TopologyEdgeDeleted(BaseEvent[EventType.TopologyEdgeDeleted]):
event_type: Literal[EventType.TopologyEdgeDeleted] = EventType.TopologyEdgeDeleted
- edge_id: TopologyEdgeId
+ edge: Connection
class TimerCreated(BaseEvent[EventType.TimerCreated]):
diff --git a/shared/types/graphs/common.py b/shared/types/graphs/common.py
deleted file mode 100644
index 301315af..00000000
--- a/shared/types/graphs/common.py
+++ /dev/null
@@ -1,172 +0,0 @@
-from collections.abc import Mapping
-from typing import Callable, Generic, Protocol, Set, Tuple, TypeVar, overload
-
-from pydantic import BaseModel
-
-from shared.types.common import NewUUID
-
-EdgeTypeT = TypeVar("EdgeTypeT", covariant=True)
-VertexTypeT = TypeVar("VertexTypeT", covariant=True)
-EdgeIdT = TypeVar("EdgeIdT", bound=NewUUID)
-VertexIdT = TypeVar("VertexIdT", bound=NewUUID)
-
-
-class VertexData(BaseModel, Generic[VertexTypeT]):
- vertex_type: VertexTypeT
-
-
-class EdgeData(BaseModel, Generic[EdgeTypeT]):
- edge_type: EdgeTypeT
-
-
-class BaseEdge(BaseModel, Generic[EdgeTypeT, EdgeIdT, VertexIdT]):
- edge_vertices: Tuple[VertexIdT, VertexIdT]
- edge_data: EdgeData[EdgeTypeT]
-
-
-class BaseVertex(BaseModel, Generic[VertexTypeT, EdgeIdT]):
- vertex_data: VertexData[VertexTypeT]
-
-
-class Vertex(
- BaseVertex[VertexTypeT, EdgeIdT], Generic[VertexTypeT, EdgeIdT, VertexIdT]
-):
- vertex_id: VertexIdT
-
-
-class Edge(
- BaseEdge[EdgeTypeT, EdgeIdT, VertexIdT], Generic[EdgeTypeT, EdgeIdT, VertexIdT]
-):
- edge_id: EdgeIdT
-
-
-class GraphData(BaseModel, Generic[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]):
- edges: Mapping[EdgeIdT, EdgeData[EdgeTypeT]] = {}
- vertices: Mapping[VertexIdT, VertexData[VertexTypeT]] = {}
-
-
-class GraphProtocol(Protocol, Generic[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]):
- def list_edges(self) -> Set[EdgeIdT]: ...
- def list_vertices(self) -> Set[VertexIdT]: ...
- def get_vertices_from_edges(
- self, edges: Set[EdgeIdT]
- ) -> Mapping[EdgeIdT, Set[VertexIdT]]: ...
- def get_edges_from_vertices(
- self, vertices: Set[VertexIdT]
- ) -> Mapping[VertexIdT, Set[EdgeIdT]]: ...
- def get_edge_data(
- self, edges: Set[EdgeIdT]
- ) -> Mapping[EdgeIdT, EdgeData[EdgeTypeT]]: ...
- def get_vertex_data(
- self, vertices: Set[VertexIdT]
- ) -> Mapping[VertexIdT, VertexData[VertexTypeT]]: ...
-
-
-class MutableGraphProtocol(GraphProtocol[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]):
- def check_edges_exists(self, edge_id: EdgeIdT) -> bool: ...
- def check_vertex_exists(self, vertex_id: VertexIdT) -> bool: ...
- def _add_edge(self, edge_id: EdgeIdT, edge_data: EdgeData[EdgeTypeT]) -> None: ...
- def _add_vertex(
- self, vertex_id: VertexIdT, vertex_data: VertexData[VertexTypeT]
- ) -> None: ...
- def _remove_edge(self, edge_id: EdgeIdT) -> None: ...
- def _remove_vertex(self, vertex_id: VertexIdT) -> None: ...
- ###
- @overload
- def attach_edge(self, edge: Edge[EdgeTypeT, EdgeIdT, VertexIdT]) -> None: ...
- @overload
- def attach_edge(
- self,
- edge: Edge[EdgeTypeT, EdgeIdT, VertexIdT],
- extra_vertex: Vertex[VertexTypeT, EdgeIdT, VertexIdT],
- ) -> None: ...
- def attach_edge(
- self,
- edge: Edge[EdgeTypeT, EdgeIdT, VertexIdT],
- extra_vertex: Vertex[VertexTypeT, EdgeIdT, VertexIdT] | None = None,
- ) -> None:
- base_vertex = edge.edge_vertices[0]
- target_vertex = edge.edge_vertices[1]
- base_vertex_exists = self.check_vertex_exists(base_vertex)
- target_vertex_exists = self.check_vertex_exists(target_vertex)
-
- if not base_vertex_exists:
- raise ValueError("Base Vertex Does Not Exist")
-
- match (target_vertex_exists, extra_vertex is not None):
- case (True, False):
- raise ValueError("New Vertex Already Exists")
- case (False, True):
- if extra_vertex is None:
- raise ValueError("BUG: Extra Vertex Must Be Provided")
- self._add_vertex(extra_vertex.vertex_id, extra_vertex.vertex_data)
- case (False, False):
- raise ValueError(
- "New Vertex Must Be Provided For Non-Existent Target Vertex"
- )
- case (True, True):
- raise ValueError("New Vertex Already Exists")
-
- self._add_edge(edge.edge_id, edge.edge_data)
-
-
-class BaseGraph(
- Generic[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
- MutableGraphProtocol[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
-):
- graph_data: GraphData[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT] = GraphData[
- EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT
- ]()
-
-
-# the first element in the return value is the filtered graph; the second is the
-# (possibly empty) set of sub-graphs that were detached during filtering.
-def filter_by_edge_data(
- graph: BaseGraph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
- keep: VertexIdT,
- predicate: Callable[[EdgeData[EdgeTypeT]], bool],
-) -> Tuple[
- BaseGraph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
- Set[BaseGraph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]],
-]: ...
-
-
-# the first element in the return value is the filtered graph; the second is the
-# (possibly empty) set of sub-graphs that were detached during filtering.
-def filter_by_vertex_data(
- graph: BaseGraph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
- keep: VertexIdT,
- predicate: Callable[[VertexData[VertexTypeT]], bool],
-) -> Tuple[
- BaseGraph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
- Set[BaseGraph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]],
-]: ...
-
-
-def map_vertices_onto_graph(
- vertices: Mapping[VertexIdT, VertexData[VertexTypeT]],
- graph: BaseGraph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
-) -> Tuple[BaseGraph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT], Set[VertexIdT]]: ...
-
-
-def map_edges_onto_graph(
- edges: Mapping[EdgeIdT, EdgeData[EdgeTypeT]],
- graph: BaseGraph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
-) -> Tuple[BaseGraph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT], Set[EdgeIdT]]: ...
-
-
-def split_graph_by_edge(
- graph: BaseGraph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
- edge: EdgeIdT,
- keep: VertexIdT,
-) -> Tuple[
- BaseGraph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
- Set[BaseGraph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]],
-]: ...
-
-
-def merge_graphs_by_edge(
- graphs: Set[BaseGraph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]],
- edge: EdgeIdT,
- keep: VertexIdT,
-) -> Tuple[BaseGraph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT], Set[EdgeIdT]]: ...
diff --git a/shared/types/graphs/topology.py b/shared/types/graphs/topology.py
deleted file mode 100644
index 75e2ecbc..00000000
--- a/shared/types/graphs/topology.py
+++ /dev/null
@@ -1,48 +0,0 @@
-from pydantic import BaseModel, IPvAnyAddress
-
-from shared.graphs import Graph
-from shared.types.common import NewUUID, NodeId
-from shared.types.profiling.common import NodePerformanceProfile
-
-
-class TopologyEdgeId(NewUUID):
- pass
-
-
-class TopologyEdgeProfile(BaseModel):
- throughput: float
- latency: float
- jitter: float
-
-
-class TopologyEdge(BaseModel):
- source_ip: IPvAnyAddress
- sink_ip: IPvAnyAddress
- edge_profile: TopologyEdgeProfile
-
-
-class TopologyNode(BaseModel):
- node_id: NodeId
- node_profile: NodePerformanceProfile
-
-
-class Topology(
- Graph[
- TopologyEdge,
- TopologyNode,
- TopologyEdgeId,
- NodeId,
- ]
-):
- pass
-
-
-class OrphanedPartOfTopology(
- Graph[
- TopologyEdge,
- TopologyNode,
- TopologyEdgeId,
- NodeId,
- ]
-):
- pass
diff --git a/shared/types/profiling.py b/shared/types/profiling.py
new file mode 100644
index 00000000..ff1af45d
--- /dev/null
+++ b/shared/types/profiling.py
@@ -0,0 +1,32 @@
+from pydantic import BaseModel, Field
+
+
+class MemoryPerformanceProfile(BaseModel):
+ ram_total: int
+ ram_used: int
+ swap_total: int
+ swap_used: int
+
+
+class SystemPerformanceProfile(BaseModel):
+ flops_fp16: float
+
+
+class NetworkInterfaceInfo(BaseModel):
+ name: str
+ ip_address: str
+ type: str
+
+
+class NodePerformanceProfile(BaseModel):
+ model_id: str
+ chip_id: str
+ memory: MemoryPerformanceProfile
+ network_interfaces: list[NetworkInterfaceInfo] = Field(default_factory=list)
+ system: SystemPerformanceProfile
+
+
+class ConnectionProfile(BaseModel):
+ throughput: float
+ latency: float
+ jitter: float
diff --git a/shared/types/profiling/common.py b/shared/types/profiling/common.py
deleted file mode 100644
index 1b318cc7..00000000
--- a/shared/types/profiling/common.py
+++ /dev/null
@@ -1,54 +0,0 @@
-from enum import Enum
-from typing import Annotated, Generic, Literal, TypeVar
-
-from pydantic import BaseModel, Field, TypeAdapter
-
-
-class ProfiledResourceName(str, Enum):
- memory = "memory"
- system = "system"
-
-
-ProfiledResourceT = TypeVar(name="ProfiledResourceT", bound=ProfiledResourceName)
-
-
-class BasePerformanceProfile(BaseModel, Generic[ProfiledResourceT]):
- """
- Details a single resource (or resource type) that is being monitored by the resource monitor.
- """
-
-
-class MemoryPerformanceProfile(BasePerformanceProfile[ProfiledResourceName.memory]):
- resource_name: Literal[ProfiledResourceName.memory] = Field(
- default=ProfiledResourceName.memory, frozen=True
- )
- ram_total: int
- ram_used: int
- swap_total: int
- swap_used: int
-
-
-class NetworkInterfaceInfo(BaseModel):
- name: str
- ip_address: str
- type: str
-
-
-class SystemPerformanceProfile(BasePerformanceProfile[ProfiledResourceName.system]):
- resource_name: Literal[ProfiledResourceName.system] = Field(
- default=ProfiledResourceName.system, frozen=True
- )
- model_id: str
- chip_id: str
- memory: int
- network_interfaces: list[NetworkInterfaceInfo] = Field(default_factory=list)
-
-
-NodePerformanceProfile = Annotated[
- MemoryPerformanceProfile | SystemPerformanceProfile,
- Field(discriminator="resource_name"),
-]
-
-NodePerformanceProfileTypeAdapter: TypeAdapter[NodePerformanceProfile] = TypeAdapter(
- NodePerformanceProfile
-)
diff --git a/shared/types/tasks/request.py b/shared/types/request.py
similarity index 100%
rename from shared/types/tasks/request.py
rename to shared/types/request.py
diff --git a/shared/types/state.py b/shared/types/state.py
index 0712d525..5851034c 100644
--- a/shared/types/state.py
+++ b/shared/types/state.py
@@ -2,17 +2,12 @@ from collections.abc import Mapping, Sequence
from enum import Enum
from typing import List
-from pydantic import BaseModel, Field, TypeAdapter
+from pydantic import BaseModel, ConfigDict, Field
+from shared.topology import Topology
from shared.types.common import NodeId
-from shared.types.graphs.topology import (
- OrphanedPartOfTopology,
- Topology,
- TopologyEdge,
- TopologyNode,
-)
-from shared.types.profiling.common import NodePerformanceProfile
-from shared.types.tasks.common import Task, TaskId, TaskSagaEntry
+from shared.types.profiling import NodePerformanceProfile
+from shared.types.tasks import Task, TaskId, TaskSagaEntry
from shared.types.worker.common import InstanceId, NodeStatus
from shared.types.worker.instances import BaseInstance
from shared.types.worker.runners import RunnerId, RunnerStatus
@@ -22,23 +17,20 @@ class ExternalCommand(BaseModel): ...
class CachePolicy(str, Enum):
- KeepAll = "KeepAll"
+ KEEP_ALL = "KEEP_ALL"
class State(BaseModel):
+ model_config = ConfigDict(arbitrary_types_allowed=True)
node_status: Mapping[NodeId, NodeStatus] = {}
instances: Mapping[InstanceId, BaseInstance] = {}
runners: Mapping[RunnerId, RunnerStatus] = {}
tasks: Mapping[TaskId, Task] = {}
task_sagas: Mapping[TaskId, Sequence[TaskSagaEntry]] = {}
node_profiles: Mapping[NodeId, NodePerformanceProfile] = {}
- topology: Topology = Topology(
- edge_base=TypeAdapter(TopologyEdge), vertex_base=TypeAdapter(TopologyNode)
- )
- history: Sequence[OrphanedPartOfTopology] = []
+ topology: Topology = Topology()
+ history: Sequence[Topology] = []
task_inbox: List[Task] = Field(default_factory=list)
task_outbox: List[Task] = Field(default_factory=list)
- cache_policy: CachePolicy = CachePolicy.KeepAll
-
- # TODO: implement / use this?
+ cache_policy: CachePolicy = CachePolicy.KEEP_ALL
last_event_applied_idx: int = Field(default=0, ge=0)
diff --git a/shared/types/tasks/common.py b/shared/types/tasks.py
similarity index 85%
rename from shared/types/tasks/common.py
rename to shared/types/tasks.py
index c324c42d..2d865f57 100644
--- a/shared/types/tasks/common.py
+++ b/shared/types/tasks.py
@@ -11,7 +11,7 @@ class TaskId(NewUUID):
pass
class TaskType(str, Enum):
- ChatCompletion = "ChatCompletion"
+ CHAT_COMPLETION = "CHAT_COMPLETION"
class TaskStatus(str, Enum):
Pending = "Pending"
@@ -22,12 +22,13 @@ class TaskStatus(str, Enum):
class Task(BaseModel):
task_id: TaskId
- task_type: TaskType # redundant atm as we only have 1 task type.
+ task_type: TaskType
instance_id: InstanceId
task_status: TaskStatus
task_params: ChatCompletionTaskParams
+
class TaskSagaEntry(BaseModel):
task_id: TaskId
instance_id: InstanceId
diff --git a/shared/types/topology.py b/shared/types/topology.py
new file mode 100644
index 00000000..ce1d97ce
--- /dev/null
+++ b/shared/types/topology.py
@@ -0,0 +1,68 @@
+from typing import Iterable, Protocol
+
+from pydantic import BaseModel, ConfigDict
+
+from shared.types.common import NewUUID, NodeId
+from shared.types.profiling import ConnectionProfile, NodePerformanceProfile
+
+
+class ConnectionId(NewUUID):
+ pass
+
+class Connection(BaseModel):
+ source_node_id: NodeId
+ sink_node_id: NodeId
+ source_multiaddr: str
+ sink_multiaddr: str
+ connection_profile: ConnectionProfile | None = None
+
+ # required for Connection to be used as a key
+ model_config = ConfigDict(frozen=True, extra="forbid", strict=True)
+ def __hash__(self) -> int:
+ return hash(
+ (
+ self.source_node_id,
+ self.sink_node_id,
+ self.source_multiaddr,
+ self.sink_multiaddr,
+ )
+ )
+ def __eq__(self, other: object) -> bool:
+ if not isinstance(other, Connection):
+ raise ValueError("Cannot compare Connection with non-Connection")
+ return (
+ self.source_node_id == other.source_node_id
+ and self.sink_node_id == other.sink_node_id
+ and self.source_multiaddr == other.source_multiaddr
+ and self.sink_multiaddr == other.sink_multiaddr
+ )
+
+
+class Node(BaseModel):
+ node_id: NodeId
+ node_profile: NodePerformanceProfile | None = None
+
+
+class TopologyProto(Protocol):
+ def add_node(self, node: Node, node_id: NodeId) -> None: ...
+
+ def add_connection(
+ self,
+ connection: Connection,
+ ) -> None: ...
+
+ def list_nodes(self) -> Iterable[Node]: ...
+
+ def list_connections(self) -> Iterable[Connection]: ...
+
+ def update_node_profile(self, node_id: NodeId, node_profile: NodePerformanceProfile) -> None: ...
+
+ def update_connection_profile(self, connection: Connection) -> None: ...
+
+ def remove_connection(self, connection: Connection) -> None: ...
+
+ def remove_node(self, node_id: NodeId) -> None: ...
+
+ def get_node_profile(self, node_id: NodeId) -> NodePerformanceProfile | None: ...
+
+ def get_connection_profile(self, connection: Connection) -> ConnectionProfile | None: ...
diff --git a/shared/types/worker/commands_runner.py b/shared/types/worker/commands_runner.py
index 7f439ddd..4432b6d7 100644
--- a/shared/types/worker/commands_runner.py
+++ b/shared/types/worker/commands_runner.py
@@ -4,7 +4,7 @@ from typing import Annotated, Generic, Literal, TypeVar
from pydantic import BaseModel, Field, TypeAdapter
from shared.openai_compat import FinishReason
-from shared.types.tasks.common import ChatCompletionTaskParams
+from shared.types.tasks import ChatCompletionTaskParams
from shared.types.worker.mlx import Host
from shared.types.worker.shards import ShardMetadata
diff --git a/shared/types/worker/ops.py b/shared/types/worker/ops.py
index 869289ff..f956a32c 100644
--- a/shared/types/worker/ops.py
+++ b/shared/types/worker/ops.py
@@ -4,7 +4,7 @@ from typing import Annotated, Generic, Literal, TypeVar, Union
from pydantic import BaseModel, Field
from shared.types.events.events import InstanceId
-from shared.types.tasks.common import Task
+from shared.types.tasks import Task
from shared.types.worker.common import RunnerId
from shared.types.worker.mlx import Host
from shared.types.worker.shards import ShardMetadata
diff --git a/shared/types/worker/resource_monitor.py b/shared/types/worker/resource_monitor.py
index f45d943a..ee5267fc 100644
--- a/shared/types/worker/resource_monitor.py
+++ b/shared/types/worker/resource_monitor.py
@@ -3,50 +3,31 @@ from abc import ABC, abstractmethod
from collections.abc import Coroutine
from typing import Callable, List, Set
-from shared.types.profiling.common import (
+from shared.types.profiling import (
MemoryPerformanceProfile,
- NodePerformanceProfile,
SystemPerformanceProfile,
)
class ResourceCollector(ABC):
- """
- Details a single resource (or resource type) that is being monitored by the resource monitor.
- """
-
- name = str
-
@abstractmethod
- async def collect(self) -> NodePerformanceProfile: ...
+ async def collect(self) -> SystemPerformanceProfile | MemoryPerformanceProfile: ...
class SystemResourceCollector(ResourceCollector):
- name = "system"
-
- @abstractmethod
async def collect(self) -> SystemPerformanceProfile: ...
class MemoryResourceCollector(ResourceCollector):
- name = "memory"
-
- @abstractmethod
async def collect(self) -> MemoryPerformanceProfile: ...
class ResourceMonitor:
data_collectors: List[ResourceCollector]
- effect_handlers: Set[Callable[[NodePerformanceProfile], None]]
-
- # Since there's no implementation, this breaks the typechecker.
- # self.collectors: list[ResourceCollector] = [
- # SystemResourceCollector(),
- # MemoryResourceCollector(),
- # ]
+ effect_handlers: Set[Callable[[SystemPerformanceProfile | MemoryPerformanceProfile], None]]
- async def _collect(self) -> list[NodePerformanceProfile]:
- tasks: list[Coroutine[None, None, NodePerformanceProfile]] = [
+ async def _collect(self) -> list[SystemPerformanceProfile | MemoryPerformanceProfile]:
+ tasks: list[Coroutine[None, None, SystemPerformanceProfile | MemoryPerformanceProfile]] = [
collector.collect() for collector in self.data_collectors
]
return await asyncio.gather(*tasks)
diff --git a/worker/runner/runner.py b/worker/runner/runner.py
index 102acfca..eebb9a5b 100644
--- a/worker/runner/runner.py
+++ b/worker/runner/runner.py
@@ -11,7 +11,7 @@ from mlx_lm.tokenizer_utils import TokenizerWrapper
from engines.mlx.utils_mlx import apply_chat_template, initialize_mlx
from shared.openai_compat import FinishReason
-from shared.types.tasks.common import ChatCompletionTaskParams
+from shared.types.tasks import ChatCompletionTaskParams
from shared.types.worker.commands_runner import (
ChatTaskMessage,
ExitMessage,
diff --git a/worker/runner/runner_supervisor.py b/worker/runner/runner_supervisor.py
index de527932..b889be4b 100644
--- a/worker/runner/runner_supervisor.py
+++ b/worker/runner/runner_supervisor.py
@@ -6,11 +6,8 @@ from types import CoroutineType
from typing import Any, Callable
from shared.types.events.chunks import GenerationChunk, TokenChunk
-from shared.types.tasks.common import (
- ChatCompletionTaskParams,
- Task,
-)
-from shared.types.tasks.request import RequestId
+from shared.types.request import RequestId
+from shared.types.tasks import ChatCompletionTaskParams, Task
from shared.types.worker.commands_runner import (
ChatTaskMessage,
ErrorResponse,
diff --git a/worker/tests/conftest.py b/worker/tests/conftest.py
index 25e226c7..b101a853 100644
--- a/worker/tests/conftest.py
+++ b/worker/tests/conftest.py
@@ -10,7 +10,7 @@ from shared.types.api import ChatCompletionMessage, ChatCompletionTaskParams
from shared.types.common import NodeId
from shared.types.models import ModelId, ModelMetadata
from shared.types.state import State
-from shared.types.tasks.common import (
+from shared.types.tasks import (
Task,
TaskId,
TaskStatus,
@@ -107,7 +107,7 @@ def completion_create_params(user_message: str) -> ChatCompletionTaskParams:
@pytest.fixture
def chat_completion_task(completion_create_params: ChatCompletionTaskParams) -> Task:
"""Creates a ChatCompletionTask directly for serdes testing"""
- return Task(task_id=TaskId(), instance_id=InstanceId(), task_type=TaskType.ChatCompletion, task_status=TaskStatus.Pending, task_params=completion_create_params)
+ return Task(task_id=TaskId(), instance_id=InstanceId(), task_type=TaskType.CHAT_COMPLETION, task_status=TaskStatus.Pending, task_params=completion_create_params)
@pytest.fixture
def chat_task(
@@ -117,7 +117,7 @@ def chat_task(
return Task(
task_id=TaskId(),
instance_id=InstanceId(),
- task_type=TaskType.ChatCompletion,
+ task_type=TaskType.CHAT_COMPLETION,
task_status=TaskStatus.Pending,
task_params=completion_create_params,
)
diff --git a/worker/tests/test_serdes.py b/worker/tests/test_serdes.py
index 7ae81bf3..6e54178b 100644
--- a/worker/tests/test_serdes.py
+++ b/worker/tests/test_serdes.py
@@ -3,7 +3,7 @@ from typing import Callable, TypeVar
from pydantic import BaseModel, TypeAdapter
-from shared.types.tasks.common import Task
+from shared.types.tasks import Task
from shared.types.worker.commands_runner import (
ChatTaskMessage,
RunnerMessageTypeAdapter,
diff --git a/worker/tests/test_supervisor.py b/worker/tests/test_supervisor.py
index 4fd1dfeb..40f4ba02 100644
--- a/worker/tests/test_supervisor.py
+++ b/worker/tests/test_supervisor.py
@@ -6,7 +6,7 @@ import pytest
from shared.openai_compat import FinishReason
from shared.types.events.chunks import TokenChunk
-from shared.types.tasks.common import (
+from shared.types.tasks import (
ChatCompletionTaskParams,
Task,
TaskType,
@@ -130,7 +130,7 @@ async def test_supervisor_early_stopping(
)
max_tokens = 50
- assert chat_task.task_type == TaskType.ChatCompletion
+ assert chat_task.task_type == TaskType.CHAT_COMPLETION
print(f'chat_task.task_params: {chat_task.task_params}')
assert isinstance(chat_task.task_params, ChatCompletionTaskParams)
task_params: ChatCompletionTaskParams = chat_task.task_params
diff --git a/worker/tests/test_worker_handlers.py b/worker/tests/test_worker_handlers.py
index d70c1ed5..0812622c 100644
--- a/worker/tests/test_worker_handlers.py
+++ b/worker/tests/test_worker_handlers.py
@@ -10,7 +10,7 @@ from shared.types.common import NodeId
from shared.types.events.chunks import TokenChunk
from shared.types.events.events import ChunkGenerated, RunnerStatusUpdated
from shared.types.events.registry import Event
-from shared.types.tasks.common import Task
+from shared.types.tasks import Task
from shared.types.worker.common import RunnerId
from shared.types.worker.instances import Instance
from shared.types.worker.ops import (
← 14b3c4a6 New API!
·
back to Exo
·
fix 76f90350 →