diff --git a/src/exo/api/main.py b/src/exo/api/main.py index fcd54c9315..eb0c360df7 100644 --- a/src/exo/api/main.py +++ b/src/exo/api/main.py @@ -89,6 +89,7 @@ PlaceInstanceParams, PlacementPreview, PlacementPreviewResponse, + PromoteMasterResponse, StartDownloadParams, StartDownloadResponse, ToolCall, @@ -167,6 +168,7 @@ ImageEdits, ImageGeneration, PlaceInstance, + PromoteMaster, SendInputChunk, SetInstanceLink, StartDownload, @@ -348,6 +350,7 @@ def _setup_routes(self) -> None: self.app.get("/instance/previews")(self.get_placement_previews) self.app.get("/instance/await", response_model=None)(self.await_instance) self.app.get("/instance/{instance_id}")(self.get_instance) + self.app.post("/master/promote/{node_id}")(self.promote_master) self.app.delete("/instance/{instance_id}")(self.delete_instance) self.app.get("/v1/instance-links")(self.list_instance_links) self.app.post("/v1/instance-links")(self.create_instance_link) @@ -692,6 +695,32 @@ async def delete_instance(self, instance_id: InstanceId) -> DeleteInstanceRespon instance_id=instance_id, ) + async def promote_master(self, node_id: NodeId) -> PromoteMasterResponse: + """Force node_id to win the next master election. + + Every node in the cluster tears down and recreates its worker + (and download coordinator) when the master changes (see + exo.main._elect_loop), so this is only allowed while the cluster is + idle -- promoting mid-serving would restart every running instance. + """ + if node_id not in set(self.state.topology.list_nodes()): + raise HTTPException(status_code=404, detail="Node not found") + if self.state.instances: + raise HTTPException( + status_code=409, + detail="Cannot promote master while instances are running " + "-- every node's worker restarts when master changes. " + "Eject running models first.", + ) + + command = PromoteMaster(target_node_id=node_id) + await self._send(command) + return PromoteMasterResponse( + message="Command received.", + command_id=command.command_id, + target_node_id=node_id, + ) + async def get_feature_flags(self) -> dict[str, bool]: return {"disaggregation": ENABLE_DISAGGREGATION} diff --git a/src/exo/api/tests/test_promote_master.py b/src/exo/api/tests/test_promote_master.py new file mode 100644 index 0000000000..c51bc0fb46 --- /dev/null +++ b/src/exo/api/tests/test_promote_master.py @@ -0,0 +1,81 @@ +# pyright: reportUnusedFunction=false, reportAny=false +from typing import Any +from unittest.mock import AsyncMock + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from exo.api.main import API +from exo.shared.models.model_cards import ModelId +from exo.shared.topology import Topology +from exo.shared.types.common import NodeId +from exo.shared.types.state import State +from exo.shared.types.worker.instances import InstanceId, MlxRingInstance +from exo.shared.types.worker.runners import ShardAssignments + +NODE_A = NodeId("node-a") +NODE_B = NodeId("node-b") + + +def _make_api(state: State) -> Any: + app = FastAPI() + api = object.__new__(API) + api.app = app + api.state = state + api._send = AsyncMock() # pyright: ignore[reportPrivateUsage] + api._setup_exception_handlers() # pyright: ignore[reportPrivateUsage] + app.post("/master/promote/{node_id}")(api.promote_master) + return api + + +def _idle_topology() -> Topology: + topology = Topology() + topology.add_node(NODE_A) + return topology + + +def test_promote_master_rejects_unknown_node() -> None: + api = _make_api(State(topology=_idle_topology(), instances={})) + client = TestClient(api.app) + + response = client.post(f"/master/promote/{NODE_B}") + + assert response.status_code == 404 + api._send.assert_not_called() + + +def test_promote_master_rejects_while_instances_running() -> None: + instance = MlxRingInstance( + instance_id=InstanceId("instance-a"), + shard_assignments=ShardAssignments( + model_id=ModelId("test-model"), runner_to_shard={}, node_to_runner={} + ), + hosts_by_node={}, + ephemeral_port=50000, + ) + api = _make_api( + State( + topology=_idle_topology(), + instances={instance.instance_id: instance}, + ) + ) + client = TestClient(api.app) + + response = client.post(f"/master/promote/{NODE_A}") + + assert response.status_code == 409 + api._send.assert_not_called() + + +def test_promote_master_sends_command_when_idle() -> None: + api = _make_api(State(topology=_idle_topology(), instances={})) + client = TestClient(api.app) + + response = client.post(f"/master/promote/{NODE_A}") + + assert response.status_code == 200 + data: dict[str, Any] = response.json() + assert data["target_node_id"] == str(NODE_A) + api._send.assert_called_once() + command = api._send.call_args[0][0] + assert command.target_node_id == NODE_A diff --git a/src/exo/api/types/__init__.py b/src/exo/api/types/__init__.py index 03f29c415e..5b4b4ee306 100644 --- a/src/exo/api/types/__init__.py +++ b/src/exo/api/types/__init__.py @@ -47,6 +47,7 @@ from .api import PlacementPreview as PlacementPreview from .api import PlacementPreviewResponse as PlacementPreviewResponse from .api import PowerUsage as PowerUsage +from .api import PromoteMasterResponse as PromoteMasterResponse from .api import PromptTokensDetails as PromptTokensDetails from .api import StartDownloadParams as StartDownloadParams from .api import StartDownloadResponse as StartDownloadResponse diff --git a/src/exo/api/types/api.py b/src/exo/api/types/api.py index 02e397adaa..aacd787657 100644 --- a/src/exo/api/types/api.py +++ b/src/exo/api/types/api.py @@ -307,6 +307,12 @@ class DeleteInstanceResponse(BaseModel): instance_id: InstanceId +class PromoteMasterResponse(BaseModel): + message: str + command_id: CommandId + target_node_id: NodeId + + class AwaitInstanceReadyMessage(BaseModel): type: Literal["ready"] = "ready" instance: Instance diff --git a/src/exo/main.py b/src/exo/main.py index f504785e6a..78302a5508 100644 --- a/src/exo/main.py +++ b/src/exo/main.py @@ -23,7 +23,7 @@ from exo.routing.event_router import EventRouter from exo.routing.router import Router, get_node_zid from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_LOG, EXO_PID_FILE -from exo.shared.election import Election, ElectionResult +from exo.shared.election import FORCE_MASTER_SENIORITY, Election, ElectionResult from exo.shared.logging import logger_cleanup, logger_setup from exo.shared.types.common import NodeId, SessionId from exo.utils import STDIO_FDS @@ -128,7 +128,7 @@ async def create(cls, args: "Args") -> Self: election = Election( node_id, # If someone manages to assemble 1 MILLION devices into an exo cluster then. well done. good job champ. - seniority=1_000_000 if args.force_master else 0, + seniority=FORCE_MASTER_SENIORITY if args.force_master else 0, # nb: this DOES feedback right now. i have thoughts on how to address this, # but ultimately it seems not worth the complexity election_message_sender=router.sender(topics.ELECTION_MESSAGES), diff --git a/src/exo/master/main.py b/src/exo/master/main.py index 485ede30f7..794e891724 100644 --- a/src/exo/master/main.py +++ b/src/exo/master/main.py @@ -28,6 +28,7 @@ ImageEdits, ImageGeneration, PlaceInstance, + PromoteMaster, RequestEventLog, SendInputChunk, SetInstanceLink, @@ -180,6 +181,11 @@ async def _command_processor(self) -> None: match command: case TestCommand(): pass + case PromoteMaster(): + # Handled by exo.shared.election.Election, which + # subscribes to the same topics.COMMANDS topic -- + # nothing for Master to do here. + pass case TextGeneration(): # set-difference => prefill-only nodes prefill_only: set[InstanceId] = set() diff --git a/src/exo/shared/election.py b/src/exo/shared/election.py index 958a83d2fa..5bd134c0b5 100644 --- a/src/exo/shared/election.py +++ b/src/exo/shared/election.py @@ -9,7 +9,7 @@ from loguru import logger from exo.routing.connection_message import ConnectionMessage -from exo.shared.types.commands import ForwarderCommand +from exo.shared.types.commands import ForwarderCommand, PromoteMaster from exo.shared.types.common import NodeId, SessionId from exo.utils.channels import Receiver, Sender from exo.utils.pydantic_ext import FrozenModel @@ -17,6 +17,12 @@ DEFAULT_ELECTION_TIMEOUT = 3.0 +# Seniority high enough that no node can realistically out-grow it through +# normal re-election (seniority only ever grows to at most the number of +# candidates seen in a round). Shared by --force-master at startup and by +# a runtime PromoteMaster command. +FORCE_MASTER_SENIORITY = 1_000_000 + class ElectionMessage(FrozenModel): clock: int @@ -83,6 +89,11 @@ def __init__( self._campaign_done: Event | None = None self._tg = TaskGroup() + # Highest seniority ever observed from a peer, used by _force_promote + # so repeated PromoteMaster commands always outrank whatever the + # cluster has seen so far, instead of tying against each other. + self._max_peer_seniority_seen = 0 + async def run(self): logger.info("Starting Election") try: @@ -132,6 +143,9 @@ async def _election_receiver(self) -> None: logger.debug("Dropping message from ourselves") # Drop messages from us (See exo.routing.router) continue + self._max_peer_seniority_seen = max( + self._max_peer_seniority_seen, message.seniority + ) # If a new round is starting, we participate if message.clock > self.clock: self.clock = message.clock @@ -181,8 +195,41 @@ async def _connection_receiver(self) -> None: async def _command_counter(self) -> None: with self._co_receiver as commands: - async for _command in commands: + async for forwarder_command in commands: self.commands_seen += 1 + command = forwarder_command.command + if ( + isinstance(command, PromoteMaster) + and command.target_node_id == self.node_id + ): + self._force_promote() + + def _force_promote(self) -> None: + """Guarantee this node wins the next election round, then trigger one. + + Sets seniority to one more than the highest value this node or any + peer has ever been observed at (floored at FORCE_MASTER_SENIORITY, the + same baseline --force-master uses at startup). Using a fixed constant + here would let a *second* PromoteMaster tie the first: both nodes + would sit at the same seniority and the round would fall through to + the commands_seen tiebreak, which favours whichever node has been + master longest -- silently no-opping the newer promotion. Always + going one higher than anything seen so far keeps repeated + promotions, and promotions away from a --force-master node, working. + """ + logger.info("Forcing this node to win the next master election") + self.seniority = ( + max( + self.seniority, + self._max_peer_seniority_seen, + FORCE_MASTER_SENIORITY - 1, + ) + + 1 + ) + self.clock += 1 + candidates: list[ElectionMessage] = [] + self._candidates = candidates + self._tg.start_soon(self._campaign, candidates, DEFAULT_ELECTION_TIMEOUT) async def _campaign( self, candidates: list[ElectionMessage], campaign_timeout: float diff --git a/src/exo/shared/tests/test_election.py b/src/exo/shared/tests/test_election.py index d1612e85e4..8cae7398cb 100644 --- a/src/exo/shared/tests/test_election.py +++ b/src/exo/shared/tests/test_election.py @@ -2,8 +2,13 @@ from anyio import create_task_group, fail_after, move_on_after from exo.routing.connection_message import ConnectionMessage -from exo.shared.election import Election, ElectionMessage, ElectionResult -from exo.shared.types.commands import ForwarderCommand, TestCommand +from exo.shared.election import ( + FORCE_MASTER_SENIORITY, + Election, + ElectionMessage, + ElectionResult, +) +from exo.shared.types.commands import ForwarderCommand, PromoteMaster, TestCommand from exo.shared.types.common import NodeId, SessionId, SystemId from exo.utils.channels import channel @@ -402,3 +407,199 @@ async def test_tie_breaker_prefers_node_with_more_commands_seen() -> None: em_in_tx.close() cm_tx.close() co_tx.close() + + +@pytest.mark.anyio +async def test_promote_master_command_forces_target_node_to_win() -> None: + """ + A peer wins a round first (so we are not master). A PromoteMaster + command targeting us should then force a new round that we win, by + boosting our own seniority the same way --force-master does at startup. + """ + em_out_tx, _em_out_rx = channel[ElectionMessage]() + em_in_tx, em_in_rx = channel[ElectionMessage]() + er_tx, er_rx = channel[ElectionResult]() + cm_tx, cm_rx = channel[ConnectionMessage]() + co_tx, co_rx = channel[ForwarderCommand]() + + me = NodeId("ME") + election = Election( + node_id=me, + election_message_receiver=em_in_rx, + election_message_sender=em_out_tx, + election_result_sender=er_tx, + connection_message_receiver=cm_rx, + command_receiver=co_rx, + is_candidate=True, + ) + + async with create_task_group() as tg: + with fail_after(2): + tg.start_soon(election.run) + + # A peer with higher seniority wins round clock=1. + await em_in_tx.send(em(clock=1, seniority=10, node_id="PEER")) + while True: + result = await er_rx.receive() + if result.session_id.election_clock == 1: + break + assert result.session_id.master_node_id == NodeId("PEER") + assert election.seniority == 0 + + # Now force-promote us. + await co_tx.send( + ForwarderCommand( + origin=SystemId("SOMEONE"), + command=PromoteMaster(target_node_id=me), + ) + ) + + # We should win the next round despite the peer's earlier win. + while True: + result = await er_rx.receive() + if result.session_id.master_node_id == me: + break + + assert election.seniority == FORCE_MASTER_SENIORITY + + em_in_tx.close() + cm_tx.close() + co_tx.close() + + +@pytest.mark.anyio +async def test_promote_master_command_for_other_node_is_ignored() -> None: + """A PromoteMaster targeting a different node must not boost our own + seniority or trigger an extra round on our side.""" + em_out_tx, _em_out_rx = channel[ElectionMessage]() + em_in_tx, em_in_rx = channel[ElectionMessage]() + er_tx, er_rx = channel[ElectionResult]() + cm_tx, cm_rx = channel[ConnectionMessage]() + co_tx, co_rx = channel[ForwarderCommand]() + + election = Election( + node_id=NodeId("ME"), + election_message_receiver=em_in_rx, + election_message_sender=em_out_tx, + election_result_sender=er_tx, + connection_message_receiver=cm_rx, + command_receiver=co_rx, + is_candidate=True, + ) + + async with create_task_group() as tg: + with fail_after(2): + tg.start_soon(election.run) + + # Consume the initial self-election result from boot (this + # naturally bumps seniority to 1 -- a solo node always wins its + # own bootstrap round against itself). + _ = await er_rx.receive() + seniority_after_boot = election.seniority + + await co_tx.send( + ForwarderCommand( + origin=SystemId("SOMEONE"), + command=PromoteMaster(target_node_id=NodeId("SOMEONE_ELSE")), + ) + ) + + # Give any (incorrect) campaign a moment to resolve, then check + # no election crowned us and our seniority wasn't force-boosted. + with move_on_after(0.3): + while True: + result = await er_rx.receive() + assert result.session_id.master_node_id != NodeId("ME") + + assert election.seniority == seniority_after_boot + + em_in_tx.close() + cm_tx.close() + co_tx.close() + + +@pytest.mark.anyio +async def test_second_promote_still_wins_over_previously_promoted_master() -> None: + """ + Regression test for a tie: boosting seniority to a fixed constant on + every PromoteMaster meant a *second* promotion tied the first on + seniority and fell to the commands_seen tiebreak, which favours + whichever node has been master longer -- silently no-opping the second + promotion. Seniority must instead be strictly higher than any previously + observed value so repeated promotions keep working. + """ + em_out_tx, em_out_rx = channel[ElectionMessage]() + em_in_tx, em_in_rx = channel[ElectionMessage]() + er_tx, er_rx = channel[ElectionResult]() + cm_tx, cm_rx = channel[ConnectionMessage]() + co_tx, co_rx = channel[ForwarderCommand]() + + me = NodeId("ME") + election = Election( + node_id=me, + election_message_receiver=em_in_rx, + election_message_sender=em_out_tx, + election_result_sender=er_tx, + connection_message_receiver=cm_rx, + command_receiver=co_rx, + is_candidate=True, + ) + + async with create_task_group() as tg: + with fail_after(2): + tg.start_soon(election.run) + + # A peer was already force-promoted earlier: it wins round 1 at + # the same fixed seniority --force-master/PromoteMaster uses, + # with a head start on commands_seen (as a long-standing master + # naturally accumulates). + await em_in_tx.send( + em( + clock=1, + seniority=FORCE_MASTER_SENIORITY, + node_id="PEER", + commands_seen=5, + ) + ) + while True: + result = await er_rx.receive() + if result.session_id.election_clock == 1: + break + assert result.session_id.master_node_id == NodeId("PEER") + + # Now promote us. Wait for our own round-2 broadcast so we know + # _force_promote has reset the candidate list before the peer's + # competing round-2 message arrives. + await co_tx.send( + ForwarderCommand( + origin=SystemId("SOMEONE"), + command=PromoteMaster(target_node_id=me), + ) + ) + while True: + got = await em_out_rx.receive() + if got.clock == 2 and got.proposed_session.master_node_id == me: + break + + # The incumbent rejoins the new round at the same fixed + # seniority it was promoted with, and with a higher + # commands_seen than we have (we've only seen the one + # PromoteMaster command so far). + await em_in_tx.send( + em( + clock=2, + seniority=FORCE_MASTER_SENIORITY, + node_id="PEER", + commands_seen=5, + ) + ) + + while True: + result = await er_rx.receive() + if result.session_id.election_clock == 2: + break + assert result.session_id.master_node_id == me + + em_in_tx.close() + cm_tx.close() + co_tx.close() diff --git a/src/exo/shared/types/commands.py b/src/exo/shared/types/commands.py index 67d318b255..f8622a8e8f 100644 --- a/src/exo/shared/types/commands.py +++ b/src/exo/shared/types/commands.py @@ -100,6 +100,16 @@ class DeleteInstanceLink(BaseCommand): link_id: InstanceLinkId +class PromoteMaster(BaseCommand): + """Force target_node_id to win the next master election. + + Delivered to every node (topics.COMMANDS is broadcast) -- only the node + whose own id matches target_node_id acts on it. + """ + + target_node_id: NodeId + + DownloadCommand = StartDownload | DeleteDownload | CancelDownload @@ -119,6 +129,7 @@ class DeleteInstanceLink(BaseCommand): | DeleteCustomModelCard | SetInstanceLink | DeleteInstanceLink + | PromoteMaster )