Add PromoteMaster command for manual master promotion

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}
 

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

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

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

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),

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()

src/exo/shared/election.py

@@ -9,14 +9,20 @@
 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
 from exo.utils.task_group import TaskGroup
 
 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