|
@@ -9,7 +9,7 @@ from . import node_service_pb2_grpc
|
|
from ..peer_handle import PeerHandle
|
|
from ..peer_handle import PeerHandle
|
|
from exo.inference.shard import Shard
|
|
from exo.inference.shard import Shard
|
|
from exo.topology.topology import Topology
|
|
from exo.topology.topology import Topology
|
|
-from exo.topology.device_capabilities import DeviceCapabilities
|
|
|
|
|
|
+from exo.topology.device_capabilities import DeviceCapabilities, DeviceFlops
|
|
from exo.helpers import DEBUG
|
|
from exo.helpers import DEBUG
|
|
|
|
|
|
|
|
|
|
@@ -117,7 +117,7 @@ class GRPCPeerHandle(PeerHandle):
|
|
response = await self.stub.CollectTopology(request)
|
|
response = await self.stub.CollectTopology(request)
|
|
topology = Topology()
|
|
topology = Topology()
|
|
for node_id, capabilities in response.nodes.items():
|
|
for node_id, capabilities in response.nodes.items():
|
|
- device_capabilities = DeviceCapabilities(model=capabilities.model, chip=capabilities.chip, memory=capabilities.memory, flops=capabilities.flops)
|
|
|
|
|
|
+ device_capabilities = DeviceCapabilities(model=capabilities.model, chip=capabilities.chip, memory=capabilities.memory, flops=DeviceFlops(fp16=capabilities.flops.fp16, fp32=capabilities.flops.fp32, int8=capabilities.flops.int8))
|
|
topology.update_node(node_id, device_capabilities)
|
|
topology.update_node(node_id, device_capabilities)
|
|
for node_id, peers in response.peer_graph.items():
|
|
for node_id, peers in response.peer_graph.items():
|
|
for peer_id in peers.peer_ids:
|
|
for peer_id in peers.peer_ids:
|