Topology update
This commit is contained in:
Seth Howes
@@ -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
|
||||
|
||||
+2
-2
@@ -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):
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
+4
-3
@@ -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],
|
||||
|
||||
@@ -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}
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
@@ -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]
|
||||
@@ -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):
|
||||
|
||||
@@ -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]):
|
||||
|
||||
@@ -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]]: ...
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
)
|
||||
+9
-17
@@ -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)
|
||||
|
||||
@@ -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
|
||||
@@ -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: ...
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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]]
|
||||
effect_handlers: Set[Callable[[SystemPerformanceProfile | MemoryPerformanceProfile], None]]
|
||||
|
||||
# Since there's no implementation, this breaks the typechecker.
|
||||
# self.collectors: list[ResourceCollector] = [
|
||||
# SystemResourceCollector(),
|
||||
# MemoryResourceCollector(),
|
||||
# ]
|
||||
|
||||
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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 (
|
||||
|
||||
Reference in New Issue
Block a user