fix: prune stale node identity state on timeout

src/exo/shared/apply.py

@@ -295,6 +295,11 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
         key: value for key, value in state.downloads.items() if key != event.node_id
     }
     # Clean up all granular node mappings
+    node_identities = {
+        key: value
+        for key, value in state.node_identities.items()
+        if key != event.node_id
+    }
     node_memory = {
         key: value for key, value in state.node_memory.items() if key != event.node_id
     }
@@ -320,6 +325,9 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
     node_rdma_ctl = {
         key: value for key, value in state.node_rdma_ctl.items() if key != event.node_id
     }
+    node_backends = {
+        key: value for key, value in state.node_backends.items() if key != event.node_id
+    }
     # Only recompute cycles if the leaving node had TB bridge enabled
     leaving_node_status = state.node_thunderbolt_bridge.get(event.node_id)
     leaving_node_had_tb_enabled = (
@@ -335,13 +343,15 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
             "downloads": downloads,
             "topology": topology,
             "last_seen": last_seen,
+            "node_identities": node_identities,
             "node_memory": node_memory,
             "node_disk": node_disk,
             "node_system": node_system,
             "node_network": node_network,
             "node_thunderbolt": node_thunderbolt,
             "node_thunderbolt_bridge": node_thunderbolt_bridge,
             "node_rdma_ctl": node_rdma_ctl,
+            "node_backends": node_backends,
             "thunderbolt_bridge_cycles": thunderbolt_bridge_cycles,
         }
     )

src/exo/shared/tests/test_apply/test_apply_node_timed_out.py

@@ -0,0 +1,41 @@
+from datetime import datetime, timezone
+
+from exo.shared.apply import apply_node_timed_out
+from exo.shared.topology import Topology
+from exo.shared.types.backends import Backend
+from exo.shared.types.common import NodeId
+from exo.shared.types.events import NodeTimedOut
+from exo.shared.types.profiling import NodeIdentity
+from exo.shared.types.state import State
+
+
+def test_apply_node_timed_out_removes_node_identity() -> None:
+    timed_out_node = NodeId("timed-out-node")
+    live_node = NodeId("live-node")
+    topology = Topology()
+    topology.add_node(timed_out_node)
+    topology.add_node(live_node)
+    live_identity = NodeIdentity(friendly_name="Live node")
+
+    state = State(
+        topology=topology,
+        last_seen={
+            timed_out_node: datetime.now(tz=timezone.utc),
+            live_node: datetime.now(tz=timezone.utc),
+        },
+        node_identities={
+            timed_out_node: NodeIdentity(friendly_name="Timed-out node"),
+            live_node: live_identity,
+        },
+        node_backends={
+            timed_out_node: [Backend.MlxMetal],
+            live_node: [Backend.MlxCpu],
+        },
+    )
+
+    new_state = apply_node_timed_out(NodeTimedOut(node_id=timed_out_node), state)
+
+    assert set(new_state.topology.list_nodes()) == {live_node}
+    assert new_state.last_seen.keys() == {live_node}
+    assert new_state.node_identities == {live_node: live_identity}
+    assert new_state.node_backends == {live_node: [Backend.MlxCpu]}