diff --git a/docs/api.md b/docs/api.md index d79524874a..9f620d0ca4 100644 --- a/docs/api.md +++ b/docs/api.md @@ -163,6 +163,43 @@ JSON describing the instance to be placed. Command acknowledgement. The instance may not be ready immediately; wait for it to appear through `/instance/await` before sending inference requests. +### Keep a Model Running + +**POST** `/deployments` + +Asks the cluster to keep a model running. Whenever the cluster has no instance +of the model (its node left, it failed to start, someone deleted it), the master +places one with the same placement `/place_instance` uses. Any instance of the +model counts, however it was placed. An instance lost before it was ready makes +the next placement wait, from 10 s doubling up to 60 s; one lost after it was +ready is placed again at once. One deployment per model. + +**Request body:** the same as `/place_instance`: + +```json +{ "model_id": "mlx-community/Llama-3.2-1B-Instruct-4bit", "sharding": "Pipeline", "instance_meta": "MlxRing", "min_nodes": 1 } +``` + +**Response:** `{ "message", "command_id", "deployment_id" }`. `409` if the +model is already kept running. + +### List Deployments + +**GET** `/deployments` + +Returns `{ "deployments": [{ "deployment": {...}, "status": ... }] }`. `status` +is `serving` (an instance of the model has every runner ready), `starting` (an +instance is on its way up), `placing` (one will be placed shortly) or +`cant_place` (no placement fits the cluster now; `deployment.placementError` +says why; it is tried again as soon as nodes or the links between them change, and every 30 s). + +### Stop Keeping a Model Running + +**DELETE** `/deployments/{deployment_id}` + +Deletes the deployment and the instance it placed. An instance of the model +placed some other way is left running. `404` if there is no such deployment. + ## 3. Models ### List Models @@ -673,6 +710,9 @@ GET /instance/placement GET /instance/{instance_id} DELETE /instance/{instance_id} POST /place_instance +GET /deployments +POST /deployments +DELETE /deployments/{deployment_id} # Models GET /models diff --git a/src/exo/api/main.py b/src/exo/api/main.py index fcd54c9315..2e6aead91c 100644 --- a/src/exo/api/main.py +++ b/src/exo/api/main.py @@ -69,6 +69,9 @@ DeleteInstanceResponse, DeleteTracesRequest, DeleteTracesResponse, + DeploymentInfo, + DeploymentList, + DeploymentResponse, ErrorInfo, ErrorResponse, FinishReason, @@ -124,6 +127,7 @@ ResponsesResponse, ) from exo.master.image_store import ImageStore +from exo.master.keeper import deployment_status from exo.master.placement import place_instance as get_instance_placements from exo.shared.apply import apply from exo.shared.constants import ( @@ -156,8 +160,10 @@ AddCustomModelCard, CancelDownload, Command, + CreateDeployment, CreateInstance, DeleteCustomModelCard, + DeleteDeployment, DeleteDownload, DeleteInstance, DeleteInstanceLink, @@ -175,6 +181,7 @@ TextGeneration, ) from exo.shared.types.common import CommandId, Id, NodeId, SystemId +from exo.shared.types.deployments import Deployment, DeploymentId from exo.shared.types.events import ( ChunkGenerated, Event, @@ -349,6 +356,9 @@ def _setup_routes(self) -> None: self.app.get("/instance/await", response_model=None)(self.await_instance) self.app.get("/instance/{instance_id}")(self.get_instance) self.app.delete("/instance/{instance_id}")(self.delete_instance) + self.app.get("/deployments")(self.list_deployments) + self.app.post("/deployments")(self.create_deployment) + self.app.delete("/deployments/{deployment_id}")(self.delete_deployment) self.app.get("/v1/instance-links")(self.list_instance_links) self.app.post("/v1/instance-links")(self.create_instance_link) self.app.put("/v1/instance-links/{link_id}")(self.update_instance_link) @@ -692,6 +702,64 @@ async def delete_instance(self, instance_id: InstanceId) -> DeleteInstanceRespon instance_id=instance_id, ) + def list_deployments(self) -> DeploymentList: + return DeploymentList( + deployments=[ + DeploymentInfo( + deployment=deployment, + status=deployment_status(deployment, self.state), + ) + for deployment in self.state.deployments.values() + ] + ) + + async def create_deployment( + self, payload: PlaceInstanceParams + ) -> DeploymentResponse: + """Keep a model running: the master places an instance of it whenever the cluster has none.""" + existing = next( + ( + deployment + for deployment in self.state.deployments.values() + if deployment.model_card.model_id == payload.model_id + ), + None, + ) + if existing is not None: + raise HTTPException( + status_code=409, + detail=f"{payload.model_id} is already kept running by deployment {existing.deployment_id}", + ) + command = CreateDeployment( + deployment=Deployment( + deployment_id=DeploymentId(), + model_card=await ModelCard.load(payload.model_id), + sharding=payload.sharding, + instance_meta=payload.instance_meta, + min_nodes=payload.min_nodes, + ) + ) + await self._send(command) + return DeploymentResponse( + message="Command received.", + command_id=command.command_id, + deployment_id=command.deployment.deployment_id, + ) + + async def delete_deployment( + self, deployment_id: DeploymentId + ) -> DeploymentResponse: + """Stop keeping a model running, and delete the instance the keeper placed for it.""" + if deployment_id not in self.state.deployments: + raise HTTPException(status_code=404, detail="Deployment not found") + command = DeleteDeployment(deployment_id=deployment_id) + await self._send(command) + return DeploymentResponse( + message="Command received.", + command_id=command.command_id, + deployment_id=deployment_id, + ) + async def get_feature_flags(self) -> dict[str, bool]: return {"disaggregation": ENABLE_DISAGGREGATION} diff --git a/src/exo/api/tests/test_deployments_api.py b/src/exo/api/tests/test_deployments_api.py new file mode 100644 index 0000000000..29c1baf12d --- /dev/null +++ b/src/exo/api/tests/test_deployments_api.py @@ -0,0 +1,144 @@ +# pyright: reportUnusedFunction=false, reportAny=false +from typing import Any +from unittest.mock import AsyncMock, patch + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from exo.api.main import API +from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask +from exo.shared.types.backends import Backend +from exo.shared.types.commands import CreateDeployment, DeleteDeployment +from exo.shared.types.deployments import Deployment, DeploymentId +from exo.shared.types.memory import Memory +from exo.shared.types.state import State +from exo.shared.types.worker.instances import InstanceMeta +from exo.shared.types.worker.shards import Sharding + +MODEL = ModelCard( + model_id=ModelId("test/model-a"), + storage_size=Memory.from_bytes(1000), + n_layers=10, + hidden_size=30, + supports_tensor=True, + tasks=[ModelTask.TextGeneration], + backends=[Backend.MlxMetal], +) + + +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.get("/deployments")(api.list_deployments) + app.post("/deployments")(api.create_deployment) + app.delete("/deployments/{deployment_id}")(api.delete_deployment) + return api + + +def _deployment() -> Deployment: + return Deployment( + deployment_id=DeploymentId(), + model_card=MODEL, + sharding=Sharding.Pipeline, + instance_meta=InstanceMeta.MlxRing, + min_nodes=1, + ) + + +def test_create_deployment_sends_the_request_to_the_master() -> None: + api = _make_api(State()) + client = TestClient(api.app) + + with patch.object(ModelCard, "load", AsyncMock(return_value=MODEL)): + response = client.post( + "/deployments", + json={ + "model_id": "test/model-a", + "sharding": "Tensor", + "instance_meta": "MlxJaccl", + "min_nodes": 2, + }, + ) + + assert response.status_code == 200 + api._send.assert_called_once() + command = api._send.call_args[0][0] + assert isinstance(command, CreateDeployment) + assert command.deployment.model_card == MODEL + assert command.deployment.sharding == Sharding.Tensor + assert command.deployment.instance_meta == InstanceMeta.MlxJaccl + assert command.deployment.min_nodes == 2 + assert response.json()["deployment_id"] == command.deployment.deployment_id + + +def test_create_deployment_refuses_a_model_already_kept_running() -> None: + existing = _deployment() + api = _make_api(State(deployments={existing.deployment_id: existing})) + client = TestClient(api.app) + + with patch.object(ModelCard, "load", AsyncMock(return_value=MODEL)): + response = client.post("/deployments", json={"model_id": "test/model-a"}) + + assert response.status_code == 409 + assert existing.deployment_id in response.json()["error"]["message"] + api._send.assert_not_called() + + +def test_delete_deployment_sends_the_request_to_the_master() -> None: + existing = _deployment() + api = _make_api(State(deployments={existing.deployment_id: existing})) + client = TestClient(api.app) + + response = client.delete(f"/deployments/{existing.deployment_id}") + + assert response.status_code == 200 + command = api._send.call_args[0][0] + assert isinstance(command, DeleteDeployment) + assert command.deployment_id == existing.deployment_id + + +def test_delete_unknown_deployment_returns_404() -> None: + api = _make_api(State()) + client = TestClient(api.app) + + response = client.delete(f"/deployments/{DeploymentId()}") + + assert response.status_code == 404 + api._send.assert_not_called() + + +def test_list_deployments_shows_each_with_its_status() -> None: + waiting = _deployment() + failing = _deployment().model_copy( + update={ + "model_card": MODEL.model_copy(update={"model_id": ModelId("test/b")}), + "placement_error": "No cycles found with sufficient memory", + } + ) + api = _make_api( + State( + deployments={ + waiting.deployment_id: waiting, + failing.deployment_id: failing, + } + ) + ) + client = TestClient(api.app) + + response = client.get("/deployments") + + assert response.status_code == 200 + listed = { + item["deployment"]["deploymentId"]: item + for item in response.json()["deployments"] + } + assert listed[waiting.deployment_id]["status"] == "placing" + assert listed[failing.deployment_id]["status"] == "cant_place" + assert ( + listed[failing.deployment_id]["deployment"]["placementError"] + == "No cycles found with sufficient memory" + ) diff --git a/src/exo/api/types/__init__.py b/src/exo/api/types/__init__.py index 03f29c415e..add5c7769a 100644 --- a/src/exo/api/types/__init__.py +++ b/src/exo/api/types/__init__.py @@ -23,6 +23,9 @@ from .api import DeleteInstanceResponse as DeleteInstanceResponse from .api import DeleteTracesRequest as DeleteTracesRequest from .api import DeleteTracesResponse as DeleteTracesResponse +from .api import DeploymentInfo as DeploymentInfo +from .api import DeploymentList as DeploymentList +from .api import DeploymentResponse as DeploymentResponse from .api import ErrorInfo as ErrorInfo from .api import ErrorResponse as ErrorResponse from .api import FinishReason as FinishReason diff --git a/src/exo/api/types/api.py b/src/exo/api/types/api.py index 02e397adaa..8b0204949c 100644 --- a/src/exo/api/types/api.py +++ b/src/exo/api/types/api.py @@ -7,6 +7,7 @@ from exo.shared.models.model_cards import ModelCard, ModelId from exo.shared.types.common import CommandId, NodeId +from exo.shared.types.deployments import Deployment, DeploymentId, DeploymentStatus from exo.shared.types.memory import Memory from exo.shared.types.text_generation import ReasoningDialect, ReasoningEffort from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta @@ -273,6 +274,21 @@ class PlaceInstanceParams(BaseModel): min_nodes: int = 1 +class DeploymentInfo(BaseModel): + deployment: Deployment + status: DeploymentStatus + + +class DeploymentList(BaseModel): + deployments: list[DeploymentInfo] + + +class DeploymentResponse(BaseModel): + message: str + command_id: CommandId + deployment_id: DeploymentId + + class CreateInstanceParams(BaseModel): instance: Instance diff --git a/src/exo/master/keeper.py b/src/exo/master/keeper.py new file mode 100644 index 0000000000..154d443c72 --- /dev/null +++ b/src/exo/master/keeper.py @@ -0,0 +1,280 @@ +"""The keeper: keeps every deployment's model running. + +Once a second the master asks the keeper what to do. For a deployment whose model has no instance +in the cluster, the keeper places one, with the same placement a `/place_instance` request uses. + +The keeper never judges or changes an existing instance. Deciding that an instance is broken stays +with the paths that delete instances today: a node going silent, runners that keep failing to start, +a user deleting it. The keeper only notices that its model has no instance any more and places a +new one. When a deployment is deleted, the keeper deletes the instance it placed for it, as +`DELETE /instance` would. + +The master's state trails the events it sends: they reach the state only after a round trip +through the event router. So the keeper remembers the instance it placed until the state shows +it, and the master emits the keeper's events and a deployment's deletion under one lock, so that +a deletion can't miss an instance that is still on its way. +""" + +import math +from dataclasses import dataclass, field + +from loguru import logger + +from exo.master.placement import ( + cancel_unnecessary_downloads, + delete_instance, + get_transition_events, + place_instance, +) +from exo.shared.types.commands import ( + DeleteDeployment, + DeleteInstance, + DownloadCommand, + PlaceInstance, +) +from exo.shared.types.common import NodeId +from exo.shared.types.deployments import Deployment, DeploymentId, DeploymentStatus +from exo.shared.types.events import ( + DeploymentDeleted, + DeploymentPlaced, + DeploymentPlacementFailed, + Event, + InstanceDeleted, +) +from exo.shared.types.state import State +from exo.shared.types.topology import RDMAConnection +from exo.shared.types.worker.instances import Instance, InstanceId +from exo.shared.types.worker.runners import RunnerReady, RunnerRunning + +# A new master waits this long before placing anything, so nodes can reconnect and report the +# instances they already run +STARTUP_GRACE = 15.0 +# How long a placement may take to appear in the state before the keeper places again +PLACEMENT_TIMEOUT = 30.0 +# How often a deployment that no placement fits is tried again; it is also tried as soon as the +# nodes placement can use, or the links between them, change +UNPLACEABLE_RETRY = 30.0 +# An instance lost before all its runners were ready counts as a failed placement: the next one +# waits FIRST_BACKOFF, doubling with each failure in a row up to MAX_BACKOFF. An instance lost +# after it was ready (its node died, someone deleted it) is placed again at once. +FIRST_BACKOFF = 10.0 +MAX_BACKOFF = 60.0 + + +# The nodes placement can use, and the links between them (source, sink, is RDMA) +_Shape = tuple[frozenset[NodeId], frozenset[tuple[NodeId, NodeId, bool]]] + + +@dataclass +class _Attempt: + """What the keeper remembers about one deployment. It is held in memory only: a new master + starts afresh, and the state tells it whether the model has an instance.""" + + instance_id: InstanceId | None = None + placed_at: float = -math.inf + ready: bool = False + failures: int = 0 + retry_at: float = -math.inf + # The cluster's shape when placement last found nothing that fits + unplaceable_on: _Shape | None = None + + +@dataclass +class Keeper: + started_at: float + _attempts: dict[DeploymentId, _Attempt] = field(default_factory=dict) + # Deployments being deleted, until the state no longer shows them + _deleted: set[DeploymentId] = field(default_factory=set) + + def step(self, state: State, now: float) -> list[Event]: + """The events that place an instance for at most one deployment that needs one. + + One placement per step, so that each placement is made against a state that already + holds the one before it. + """ + for deployment_id in list(self._attempts): + if deployment_id not in state.deployments: + del self._attempts[deployment_id] + self._deleted.intersection_update(state.deployments) + if now - self.started_at < STARTUP_GRACE: + return [] + for deployment in sorted( + state.deployments.values(), key=lambda d: d.deployment_id + ): + if deployment.deployment_id in self._deleted: + continue + attempt = self._attempts.setdefault(deployment.deployment_id, _Attempt()) + _observe(deployment, attempt, state, now) + if events := _keep(deployment, attempt, state, now): + return events + return [] + + def delete( + self, command: DeleteDeployment, state: State + ) -> tuple[list[Event], list[DownloadCommand]]: + """Stop keeping a model running, and delete the instance the keeper placed for it, exactly + as `DELETE /instance` does, even if that instance hasn't reached the state yet. An instance + someone else placed for the model is left alone.""" + deployment = state.deployments.get(command.deployment_id) + if deployment is None: + return [], [] + self._deleted.add(deployment.deployment_id) + events: list[Event] = [ + DeploymentDeleted(deployment_id=deployment.deployment_id) + ] + attempt = self._attempts.pop(deployment.deployment_id, None) + if ( + attempt is not None + and attempt.instance_id is not None + and not _arrived(deployment, attempt) + ): + # Its InstanceCreated was sent before this, so the deletion is applied after it + events.append(InstanceDeleted(instance_id=attempt.instance_id)) + instance_id = deployment.instance_id + if instance_id is None or instance_id not in state.instances: + return events, [] + placement = delete_instance( + DeleteInstance(instance_id=instance_id), state.instances + ) + events.extend(get_transition_events(state.instances, placement, state.tasks)) + return events, list(cancel_unnecessary_downloads(placement, state.downloads)) + + +def _observe( + deployment: Deployment, attempt: _Attempt, state: State, now: float +) -> None: + """Follow the instance the keeper placed last: on its way, there, ready, or lost.""" + if attempt.instance_id is None or not _arrived(deployment, attempt): + return + instance = state.instances.get(attempt.instance_id) + if instance is not None: + if is_ready(instance, state): + attempt.ready = True + attempt.failures = 0 + return + if not attempt.ready: + attempt.failures += 1 + attempt.retry_at = now + min( + FIRST_BACKOFF * 2.0 ** (attempt.failures - 1), MAX_BACKOFF + ) + attempt.instance_id = None + attempt.ready = False + + +def _arrived(deployment: Deployment, attempt: _Attempt) -> bool: + """Whether the keeper's last placement has reached the state. DeploymentPlaced follows the + instance's InstanceCreated, so once the deployment names the instance, the instance has been + created, even if it is already gone again.""" + return deployment.instance_id == attempt.instance_id + + +def _keep( + deployment: Deployment, attempt: _Attempt, state: State, now: float +) -> list[Event]: + if has_instance(deployment, state): + return [] + in_flight = ( + attempt.instance_id is not None + and not _arrived(deployment, attempt) + and now - attempt.placed_at < PLACEMENT_TIMEOUT + ) + shape = _shape(state) + cluster_changed = attempt.unplaceable_on not in (None, shape) + if in_flight or (now < attempt.retry_at and not cluster_changed): + return [] + command = PlaceInstance( + model_card=deployment.model_card, + sharding=deployment.sharding, + instance_meta=deployment.instance_meta, + min_nodes=deployment.min_nodes, + ) + try: + placement = place_instance( + command, + state.topology, + state.instances, + state.node_memory, + state.node_network, + state.node_backends, + download_status=state.downloads, + node_rdma_ctl=state.node_rdma_ctl, + ) + (instance_id,) = [i for i in placement if i not in state.instances] + except Exception as error: + # A model that fits nowhere says why with a ValueError. Anything else placement raises + # is this deployment's problem too: it must not stop the keeper keeping the others. + attempt.retry_at = now + UNPLACEABLE_RETRY + attempt.unplaceable_on = shape + reason = ( + str(error) + if isinstance(error, ValueError) + else f"{type(error).__name__}: {error}" + ) + if deployment.placement_error == reason: + return [] + if not isinstance(error, ValueError): + logger.opt(exception=error).warning( + f"Placing {deployment.model_card.model_id} failed unexpectedly" + ) + return [ + DeploymentPlacementFailed( + deployment_id=deployment.deployment_id, error=reason + ) + ] + attempt.instance_id, attempt.placed_at, attempt.ready = instance_id, now, False + attempt.unplaceable_on = None + return [ + *get_transition_events(state.instances, placement, state.tasks), + DeploymentPlaced( + deployment_id=deployment.deployment_id, instance_id=instance_id + ), + ] + + +def _shape(state: State) -> _Shape: + """What placement works from, but for memory figures, which change all the time. A joining + node arrives in pieces: it shows up, reports its memory and backends, and its links appear + seconds later, so the keeper tries again as each piece arrives.""" + nodes = frozenset( + node_id + for node_id in state.topology.list_nodes() + if node_id in state.node_memory and node_id in state.node_backends + ) + links = frozenset( + (link.source, link.sink, isinstance(link.edge, RDMAConnection)) + for link in state.topology.list_connections() + if link.source in nodes and link.sink in nodes + ) + return nodes, links + + +def has_instance(deployment: Deployment, state: State) -> bool: + """Whether the cluster has an instance of the deployment's model, however it was created.""" + return any( + instance.shard_assignments.model_id == deployment.model_card.model_id + for instance in state.instances.values() + ) + + +def is_ready(instance: Instance, state: State) -> bool: + """Whether every runner of the instance has loaded its shard and can serve.""" + runner_ids = list(instance.shard_assignments.node_to_runner.values()) + return bool(runner_ids) and all( + isinstance(state.runners.get(runner_id), (RunnerReady, RunnerRunning)) + for runner_id in runner_ids + ) + + +def deployment_status(deployment: Deployment, state: State) -> DeploymentStatus: + """What a deployment is doing, for the API: serving when an instance of its model is ready, + starting while one is on its way up, otherwise placing or unable to place.""" + instances = [ + instance + for instance in state.instances.values() + if instance.shard_assignments.model_id == deployment.model_card.model_id + ] + if any(is_ready(instance, state) for instance in instances): + return "serving" + if instances: + return "starting" + return "cant_place" if deployment.placement_error is not None else "placing" diff --git a/src/exo/master/main.py b/src/exo/master/main.py index 485ede30f7..a05721e593 100644 --- a/src/exo/master/main.py +++ b/src/exo/master/main.py @@ -1,8 +1,10 @@ from datetime import datetime, timedelta, timezone +from time import monotonic import anyio from loguru import logger +from exo.master.keeper import Keeper from exo.master.placement import ( add_instance_to_placements, cancel_unnecessary_downloads, @@ -19,8 +21,10 @@ from exo.shared.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED from exo.shared.types.commands import ( AddCustomModelCard, + CreateDeployment, CreateInstance, DeleteCustomModelCard, + DeleteDeployment, DeleteInstance, DeleteInstanceLink, ForwarderCommand, @@ -40,6 +44,7 @@ from exo.shared.types.events import ( CustomModelCardAdded, CustomModelCardDeleted, + DeploymentCreated, Event, GlobalForwarderEvent, IndexedEvent, @@ -146,6 +151,9 @@ def __init__( self._event_log = DiskEventLog(EXO_EVENT_LOG_DIR / "master") self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {} self._expected_ranks: dict[TaskId, set[int]] = {} + self._keeper = Keeper(started_at=monotonic()) + # Held while the keeper's events, or a deployment's deletion, are worked out and sent + self._keeping = anyio.Lock() async def run(self): logger.info("Starting Master") @@ -155,6 +163,7 @@ async def run(self): tg.start_soon(self._event_processor) tg.start_soon(self._command_processor) tg.start_soon(self._plan) + tg.start_soon(self._keep_deployments) except* (EventRouterBrokenResourceError, EventRouterClosedResourceError): # Event router has been closed (try-star syntax handles error groups) pass @@ -436,6 +445,23 @@ async def _command_processor(self) -> None: generated_events.append( CustomModelCardDeleted(model_id=command.model_id) ) + case CreateDeployment(): + generated_events.append( + DeploymentCreated(deployment=command.deployment) + ) + case DeleteDeployment(): + async with self._keeping: + events, download_commands = self._keeper.delete( + command, self.state + ) + for event in events: + await self.event_sender.send(event) + for cmd in download_commands: + await self.download_command_sender.send( + ForwarderDownloadCommand( + origin=self._system_id, command=cmd + ) + ) case SetInstanceLink(): link = InstanceLink( link_id=command.link_id, @@ -467,6 +493,17 @@ async def _command_processor(self) -> None: except Exception as e: logger.opt(exception=e).warning("Error in command processor") + async def _keep_deployments(self) -> None: + """Once a second, place an instance for a deployment whose model has none (see keeper.py).""" + while True: + await anyio.sleep(1) + try: + async with self._keeping: + for event in self._keeper.step(self.state, monotonic()): + await self.event_sender.send(event) + except Exception as e: + logger.opt(exception=e).warning("Error keeping deployments running") + # These plan loops are the cracks showing in our event sourcing architecture - more things could be commands async def _plan(self) -> None: while True: diff --git a/src/exo/master/tests/test_keeper.py b/src/exo/master/tests/test_keeper.py new file mode 100644 index 0000000000..fde9f90ad3 --- /dev/null +++ b/src/exo/master/tests/test_keeper.py @@ -0,0 +1,982 @@ +"""The keeper as a pure function of the cluster state. + +Each test builds a State through the same events the cluster applies, steps a Keeper at chosen +times, and applies what it returns, so the tests follow the keeper exactly as the master runs it. +""" + +import copy +import random +from collections.abc import Callable, Iterable + +import pytest + +from exo.master import keeper as keeper_module +from exo.master.keeper import ( + FIRST_BACKOFF, + MAX_BACKOFF, + PLACEMENT_TIMEOUT, + STARTUP_GRACE, + UNPLACEABLE_RETRY, + Keeper, + deployment_status, + has_instance, +) +from exo.master.placement import place_instance +from exo.master.tests.conftest import ( + create_node_memory, + create_node_network, + create_socket_connection, +) +from exo.shared.apply import event_apply +from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask +from exo.shared.topology import Topology +from exo.shared.types.backends import Backend +from exo.shared.types.commands import ( + CancelDownload, + DeleteDeployment, + PlaceInstance, +) +from exo.shared.types.common import NodeId +from exo.shared.types.deployments import Deployment, DeploymentId +from exo.shared.types.events import ( + DeploymentCreated, + DeploymentDeleted, + DeploymentPlaced, + DeploymentPlacementFailed, + Event, + InstanceCreated, + InstanceDeleted, + NodeTimedOut, + RunnerStatusUpdated, +) +from exo.shared.types.memory import Memory +from exo.shared.types.state import State +from exo.shared.types.topology import Connection +from exo.shared.types.worker.downloads import DownloadOngoing, DownloadProgressData +from exo.shared.types.worker.instances import InstanceId, InstanceMeta +from exo.shared.types.worker.runners import ( + RunnerFailed, + RunnerLoading, + RunnerReady, + RunnerRunning, +) +from exo.shared.types.worker.shards import Sharding + +START = 1000.0 +AFTER_GRACE = START + STARTUP_GRACE +NODE_MEMORY = 10_000_000 + + +def _model(name: str, size_bytes: int = 1_000_000) -> ModelCard: + return ModelCard( + model_id=ModelId(name), + storage_size=Memory.from_bytes(size_bytes), + n_layers=10, + hidden_size=30, + supports_tensor=True, + tasks=[ModelTask.TextGeneration], + backends=[Backend.MlxMetal], + ) + + +MODEL_A = _model("test/model-a") +MODEL_B = _model("test/model-b") + + +def _cluster(node_count: int, memory: int = NODE_MEMORY) -> State: + """A cluster of fully connected Metal nodes.""" + nodes = [NodeId() for _ in range(node_count)] + topology = Topology() + for node in nodes: + topology.add_node(node) + port = 0 + for source in nodes: + for sink in nodes: + if source != sink: + port += 1 + topology.add_connection( + Connection( + source=source, sink=sink, edge=create_socket_connection(port) + ) + ) + return State( + topology=topology, + node_memory={node: create_node_memory(memory) for node in nodes}, + node_network={node: create_node_network() for node in nodes}, + node_backends={node: [Backend.MlxMetal] for node in nodes}, + ) + + +def _deployment(model: ModelCard) -> Deployment: + return Deployment( + deployment_id=DeploymentId(), + model_card=model, + sharding=Sharding.Pipeline, + instance_meta=InstanceMeta.MlxRing, + min_nodes=1, + ) + + +def _apply(state: State, events: Iterable[Event]) -> State: + for event in events: + state = event_apply(event, state) + return state + + +def _deploy(state: State, *models: ModelCard) -> tuple[State, list[Deployment]]: + deployments = [_deployment(model) for model in models] + state = _apply(state, [DeploymentCreated(deployment=d) for d in deployments]) + return state, deployments + + +def _without_ids(events: Iterable[Event]) -> list[tuple[str, dict[str, object]]]: + """Events to compare: every event gets a fresh event_id.""" + return [(type(e).__name__, e.model_dump(exclude={"event_id"})) for e in events] + + +def _placed(events: list[Event]) -> InstanceId: + """The instance a keeper step placed.""" + (placed,) = [e for e in events if isinstance(e, DeploymentPlaced)] + return placed.instance_id + + +def _make_ready(state: State, instance_id: InstanceId) -> State: + runner_ids = state.instances[instance_id].shard_assignments.node_to_runner.values() + return _apply( + state, + [ + RunnerStatusUpdated(runner_id=runner_id, runner_status=RunnerReady()) + for runner_id in runner_ids + ], + ) + + +def _lose(state: State, instance_id: InstanceId) -> State: + return event_apply(InstanceDeleted(instance_id=instance_id), state) + + +def _instances_of(state: State, model: ModelCard) -> list[InstanceId]: + return [ + instance_id + for instance_id, instance in state.instances.items() + if instance.shard_assignments.model_id == model.model_id + ] + + +def _place_by_hand(state: State, model: ModelCard) -> tuple[State, InstanceId]: + """Place an instance as a /place_instance request would, without the keeper.""" + placement = place_instance( + PlaceInstance( + model_card=model, + sharding=Sharding.Pipeline, + instance_meta=InstanceMeta.MlxRing, + min_nodes=1, + ), + state.topology, + state.instances, + state.node_memory, + state.node_network, + state.node_backends, + ) + (instance_id,) = [i for i in placement if i not in state.instances] + return ( + event_apply(InstanceCreated(instance=placement[instance_id]), state), + instance_id, + ) + + +# Placing + + +def test_does_nothing_without_deployments(): + keeper = Keeper(started_at=START) + state = _cluster(2) + + assert keeper.step(state, AFTER_GRACE) == [] + assert keeper.step(state, AFTER_GRACE + 3600) == [] + + +def test_waits_for_nodes_to_report_their_instances_after_starting(): + keeper = Keeper(started_at=START) + state, _ = _deploy(_cluster(1), MODEL_A) + + for now in (START, START + 1, AFTER_GRACE - 0.001): + assert keeper.step(state, now) == [] + assert keeper.step(state, AFTER_GRACE) != [] + + +def test_places_an_instance_for_a_deployment_whose_model_has_none(): + keeper = Keeper(started_at=START) + state, (deployment,) = _deploy(_cluster(2), MODEL_A) + + events = keeper.step(state, AFTER_GRACE) + state = _apply(state, events) + + instance_id = _placed(events) + created = [e for e in events if isinstance(e, InstanceCreated)] + assert [e.instance.instance_id for e in created] == [instance_id] + assert state.instances[instance_id].shard_assignments.model_id == MODEL_A.model_id + assert state.deployments[deployment.deployment_id].instance_id == instance_id + assert keeper.step(state, AFTER_GRACE + 1) == [] + + +def test_places_with_the_deployments_options(): + keeper = Keeper(started_at=START) + deployment = _deployment(MODEL_A).model_copy(update={"min_nodes": 3}) + state = event_apply(DeploymentCreated(deployment=deployment), _cluster(3)) + + instance_id = _placed(events := keeper.step(state, AFTER_GRACE)) + state = _apply(state, events) + + assert len(state.instances[instance_id].shard_assignments.node_to_runner) == 3 + + +def test_any_instance_of_the_model_keeps_the_deployment_satisfied(): + keeper = Keeper(started_at=START) + state, _ = _deploy(_cluster(2), MODEL_A) + state, _ = _place_by_hand(state, MODEL_A) + + for now in (AFTER_GRACE, AFTER_GRACE + 60, AFTER_GRACE + 3600): + assert keeper.step(state, now) == [] + + +def test_an_instance_of_another_model_does_not_count(): + keeper = Keeper(started_at=START) + state, _ = _deploy(_cluster(2), MODEL_A) + state, _ = _place_by_hand(state, MODEL_B) + + instance_id = _placed(events := keeper.step(state, AFTER_GRACE)) + state = _apply(state, events) + + assert state.instances[instance_id].shard_assignments.model_id == MODEL_A.model_id + assert len(_instances_of(state, MODEL_B)) == 1 + + +def test_places_for_one_deployment_per_step(): + keeper = Keeper(started_at=START) + state, _ = _deploy(_cluster(2), MODEL_A, MODEL_B) + + first = keeper.step(state, AFTER_GRACE) + assert len([e for e in first if isinstance(e, InstanceCreated)]) == 1 + state = _apply(state, first) + + second = keeper.step(state, AFTER_GRACE + 1) + assert len([e for e in second if isinstance(e, InstanceCreated)]) == 1 + state = _apply(state, second) + + assert len(_instances_of(state, MODEL_A)) == 1 + assert len(_instances_of(state, MODEL_B)) == 1 + assert keeper.step(state, AFTER_GRACE + 2) == [] + + +def test_only_ever_creates_instances(): + """A step adds one instance; it never deletes or changes one that exists.""" + keeper = Keeper(started_at=START) + state, _ = _deploy(_cluster(3), MODEL_A, MODEL_B) + state, _ = _place_by_hand(state, _model("test/someone-elses")) + + events = keeper.step(state, AFTER_GRACE) + + assert {type(e) for e in events} == {InstanceCreated, DeploymentPlaced} + after = _apply(state, events) + for instance_id, instance in state.instances.items(): + assert after.instances[instance_id] == instance + + +# Waiting for a placement to show up + + +def test_does_not_place_again_while_the_last_placement_is_on_its_way(): + keeper = Keeper(started_at=START) + state, _ = _deploy(_cluster(2), MODEL_A) + + first = _placed(keeper.step(state, AFTER_GRACE)) + # The events haven't reached the state yet + for now in (AFTER_GRACE + 1, AFTER_GRACE + PLACEMENT_TIMEOUT - 0.001): + assert keeper.step(state, now) == [] + + second = _placed(keeper.step(state, AFTER_GRACE + PLACEMENT_TIMEOUT)) + assert second != first + + +def test_a_placement_that_never_arrives_is_not_counted_as_a_failure(): + keeper = Keeper(started_at=START) + state, _ = _deploy(_cluster(2), MODEL_A) + + keeper.step(state, AFTER_GRACE) + now = AFTER_GRACE + PLACEMENT_TIMEOUT + events = keeper.step(state, now) + state = _apply(state, events) + + assert has_instance(next(iter(state.deployments.values())), state) + + +# Losing an instance + + +def test_places_again_at_once_when_a_ready_instance_is_lost(): + keeper = Keeper(started_at=START) + state, _ = _deploy(_cluster(2), MODEL_A) + now = AFTER_GRACE + for _ in range(5): + instance_id = _placed(events := keeper.step(state, now)) + state = _make_ready(_apply(state, events), instance_id) + now += 1 + assert keeper.step(state, now) == [] + state = _lose(state, instance_id) + now += 1 + + +def test_places_again_when_someone_deletes_the_only_instance(): + keeper = Keeper(started_at=START) + state, _ = _deploy(_cluster(2), MODEL_A) + state, by_hand = _place_by_hand(state, MODEL_A) + assert keeper.step(state, AFTER_GRACE) == [] + + state = _lose(state, by_hand) + + assert _placed(keeper.step(state, AFTER_GRACE + 1)) != by_hand + + +def test_places_again_on_the_nodes_left_when_a_node_dies(): + keeper = Keeper(started_at=START) + state, _ = _deploy(_cluster(2), MODEL_A) + instance_id = _placed(events := keeper.step(state, AFTER_GRACE)) + state = _make_ready(_apply(state, events), instance_id) + assert keeper.step(state, AFTER_GRACE + 1) == [] + (dead,) = state.instances[instance_id].shard_assignments.node_to_runner + + state = _apply( + state, [NodeTimedOut(node_id=dead), InstanceDeleted(instance_id=instance_id)] + ) + events = keeper.step(state, AFTER_GRACE + 2) + state = _apply(state, events) + + (replacement,) = _instances_of(state, MODEL_A) + assert dead not in state.instances[replacement].shard_assignments.node_to_runner + + +def test_an_instance_lost_between_two_steps_is_placed_again_after_one_backoff(): + """The keeper never sees this instance: it arrives and is gone before the next step. It still + counts as placed and lost, so the keeper doesn't wait out PLACEMENT_TIMEOUT for it.""" + keeper = Keeper(started_at=START) + state, _ = _deploy(_cluster(2), MODEL_A) + instance_id = _placed(events := keeper.step(state, AFTER_GRACE)) + state = _lose(_apply(state, events), instance_id) + + assert keeper.step(state, AFTER_GRACE + 1) == [] + assert FIRST_BACKOFF < PLACEMENT_TIMEOUT + assert keeper.step(state, AFTER_GRACE + 1 + FIRST_BACKOFF) != [] + + +def test_backs_off_when_instances_are_lost_before_they_are_ready(): + keeper = Keeper(started_at=START) + state, _ = _deploy(_cluster(2), MODEL_A) + now = AFTER_GRACE + # 10 s, 20 s, 40 s, then the 60 s cap + waits = [min(FIRST_BACKOFF * 2.0**i, MAX_BACKOFF) for i in range(6)] + assert waits[-3:] == [MAX_BACKOFF] * 3 + + for wait in waits: + instance_id = _placed(events := keeper.step(state, now)) + state = _apply(state, events) + assert keeper.step(state, now + 1) == [] # still loading + state = _lose(state, instance_id) + lost_at = now + 2 + assert keeper.step(state, lost_at) == [] + assert keeper.step(state, lost_at + wait - 0.001) == [] + now = lost_at + wait + + +def test_a_ready_instance_clears_the_backoff(): + keeper = Keeper(started_at=START) + state, _ = _deploy(_cluster(2), MODEL_A) + + def place_then_lose(now: float, ready: bool) -> None: + """Place an instance at `now`, see it a second later, then lose it.""" + nonlocal state + instance_id = _placed(events := keeper.step(state, now)) + state = _apply(state, events) + if ready: + state = _make_ready(state, instance_id) + assert keeper.step(state, now + 1) == [] + state = _lose(state, instance_id) + + first = AFTER_GRACE + place_then_lose(first, ready=False) + assert keeper.step(state, first + 2) == [] # waits FIRST_BACKOFF + second = first + 2 + FIRST_BACKOFF + place_then_lose(second, ready=False) + assert keeper.step(state, second + 2) == [] # waits twice that + third = second + 2 + 2 * FIRST_BACKOFF + place_then_lose(third, ready=True) + + # Lost after it was ready: placed again at once, and a failure after that waits only + # FIRST_BACKOFF again + place_then_lose(third + 2, ready=False) + lost_at = third + 4 + assert keeper.step(state, lost_at) == [] + assert keeper.step(state, lost_at + FIRST_BACKOFF - 0.001) == [] + assert keeper.step(state, lost_at + FIRST_BACKOFF) != [] + + +def test_a_runner_that_becomes_ready_late_still_counts(): + keeper = Keeper(started_at=START) + state, _ = _deploy(_cluster(1), MODEL_A) + instance_id = _placed(events := keeper.step(state, AFTER_GRACE)) + state = _apply(state, events) + (runner_id,) = state.instances[ + instance_id + ].shard_assignments.node_to_runner.values() + + state = event_apply( + RunnerStatusUpdated(runner_id=runner_id, runner_status=RunnerLoading()), state + ) + keeper.step(state, AFTER_GRACE + 1) + state = event_apply( + RunnerStatusUpdated(runner_id=runner_id, runner_status=RunnerRunning()), state + ) + keeper.step(state, AFTER_GRACE + 60) + state = _lose(state, instance_id) + + assert keeper.step(state, AFTER_GRACE + 61) != [] + + +# A deployment no placement fits + + +def test_reports_once_why_a_deployment_cant_be_placed_and_keeps_trying(): + keeper = Keeper(started_at=START) + too_big = _model("test/too-big", size_bytes=NODE_MEMORY * 10) + state, (deployment,) = _deploy(_cluster(2), too_big) + + events = keeper.step(state, AFTER_GRACE) + (failed,) = events + assert isinstance(failed, DeploymentPlacementFailed) + state = _apply(state, events) + assert state.deployments[deployment.deployment_id].placement_error == failed.error + + assert keeper.step(state, AFTER_GRACE + 1) == [] + # Tried again, failed the same way: nothing new to say + assert keeper.step(state, AFTER_GRACE + UNPLACEABLE_RETRY) == [] + assert keeper.step(state, AFTER_GRACE + 10 * UNPLACEABLE_RETRY) == [] + + +def test_reports_when_the_reason_it_cant_be_placed_changes(): + keeper = Keeper(started_at=START) + too_big = _model("test/too-big", size_bytes=NODE_MEMORY * 10) + state, _ = _deploy(_cluster(2), too_big) + state = _apply(state, keeper.step(state, AFTER_GRACE)) + + roomy = _cluster(2, memory=NODE_MEMORY * 100) + state = state.model_copy( + update={ + "topology": roomy.topology, + "node_memory": roomy.node_memory, + "node_network": roomy.node_network, + "node_backends": {}, + } + ) + (failed,) = keeper.step(state, AFTER_GRACE + UNPLACEABLE_RETRY) + + assert isinstance(failed, DeploymentPlacementFailed) + assert "backend" in failed.error + + +def test_places_once_memory_frees_up_on_the_same_nodes(): + keeper = Keeper(started_at=START) + big = _model("test/big", size_bytes=int(NODE_MEMORY * 2.5)) + state, (deployment,) = _deploy(_cluster(3, memory=NODE_MEMORY // 2), big) + state = _apply(state, keeper.step(state, AFTER_GRACE)) + assert ( + deployment_status(state.deployments[deployment.deployment_id], state) + == "cant_place" + ) + + state = state.model_copy( + update={ + "node_memory": { + node_id: create_node_memory(NODE_MEMORY) + for node_id in state.node_memory + } + } + ) + # The same nodes: tried again on the periodic retry + assert keeper.step(state, AFTER_GRACE + UNPLACEABLE_RETRY - 0.001) == [] + state = _apply(state, keeper.step(state, AFTER_GRACE + UNPLACEABLE_RETRY)) + + assert len(_instances_of(state, big)) == 1 + assert state.deployments[deployment.deployment_id].placement_error is None + + +def test_tries_again_at_once_when_a_node_joins(): + keeper = Keeper(started_at=START) + big = _model("test/big", size_bytes=int(NODE_MEMORY * 2.5)) + full = _cluster(3) + joining = next(iter(full.node_memory)) + state, _ = _deploy(event_apply(NodeTimedOut(node_id=joining), full), big) + state = _apply(state, keeper.step(state, AFTER_GRACE)) + assert _instances_of(state, big) == [] + + state = state.model_copy( + update={ + "topology": full.topology, + "node_memory": full.node_memory, + "node_backends": full.node_backends, + } + ) + state = _apply(state, keeper.step(state, AFTER_GRACE + 2)) + + assert len(_instances_of(state, big)) == 1 + + +def test_a_joining_node_is_tried_again_as_each_piece_of_it_arrives(): + """A node joins in pieces: it shows up, reports its memory and backends, and its links come + seconds later. Each piece is a reason to try again; otherwise the try made after the first + piece fails, and the keeper waits a full UNPLACEABLE_RETRY.""" + keeper = Keeper(started_at=START) + big = _model("test/big", size_bytes=int(NODE_MEMORY * 2.5)) + full = _cluster(3) + joining = next(iter(full.node_memory)) + state, _ = _deploy(event_apply(NodeTimedOut(node_id=joining), full), big) + state = _apply(state, keeper.step(state, AFTER_GRACE)) + + shown_up = copy.deepcopy(state.topology) + shown_up.add_node(joining) + state = state.model_copy(update={"topology": shown_up}) + assert keeper.step(state, AFTER_GRACE + 1) == [] + state = state.model_copy( + update={"node_memory": full.node_memory, "node_backends": full.node_backends} + ) + assert keeper.step(state, AFTER_GRACE + 2) == [] # no links to it yet + state = state.model_copy(update={"topology": full.topology}) + state = _apply(state, keeper.step(state, AFTER_GRACE + 3)) + + assert len(_instances_of(state, big)) == 1 + + +def test_an_unplaceable_deployment_does_not_hold_up_the_others(): + keeper = Keeper(started_at=START) + too_big = _model("test/too-big", size_bytes=NODE_MEMORY * 10) + state, _ = _deploy(_cluster(2), too_big, MODEL_A) + + now = AFTER_GRACE + for _ in range(3): + state = _apply(state, keeper.step(state, now)) + now += 1 + + assert len(_instances_of(state, MODEL_A)) == 1 + + +def _crashing_for[**P, R](model_id: ModelId, place: Callable[P, R]) -> Callable[P, R]: + """`place`, but raising an unexpected error for one model.""" + + def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: + command = args[0] + if ( + isinstance(command, PlaceInstance) + and command.model_card.model_id == model_id + ): + raise KeyError("some node") + return place(*args, **kwargs) + + return wrapper + + +def test_a_deployment_whose_placement_crashes_does_not_hold_up_the_others( + monkeypatch: pytest.MonkeyPatch, +): + """Placement raising something unexpected for one model is that deployment's problem: it is + reported and tried again later, and the other deployments are still kept.""" + broken = _model("test/broken") + monkeypatch.setattr( + keeper_module, + "place_instance", + _crashing_for(broken.model_id, keeper_module.place_instance), + ) + keeper = Keeper(started_at=START) + # The keeper goes through deployments in id order: the broken one first + failing = _deployment(broken).model_copy( + update={"deployment_id": DeploymentId("a")} + ) + working = _deployment(MODEL_A).model_copy( + update={"deployment_id": DeploymentId("b")} + ) + state = _apply( + _cluster(2), + [DeploymentCreated(deployment=failing), DeploymentCreated(deployment=working)], + ) + + events = keeper.step(state, AFTER_GRACE) + (failed,) = events + assert isinstance(failed, DeploymentPlacementFailed) + assert failed.error == "KeyError: 'some node'" + state = _apply(state, events) + state = _apply(state, keeper.step(state, AFTER_GRACE + 1)) + + assert len(_instances_of(state, MODEL_A)) == 1 + assert keeper.step(state, AFTER_GRACE + 2) == [] + assert ( + deployment_status(state.deployments[failing.deployment_id], state) + == "cant_place" + ) + + +# Deployments coming and going + + +def test_forgets_a_deployment_that_leaves_the_state_without_it(): + """As when another master deleted it, or the state was replaced by a newer snapshot.""" + keeper = Keeper(started_at=START) + state, (deployment,) = _deploy(_cluster(2), MODEL_A) + instance_id = _placed(events := keeper.step(state, AFTER_GRACE)) + state = _apply(state, events) + keeper.step(state, AFTER_GRACE + 1) + state = _lose(state, instance_id) + assert keeper.step(state, AFTER_GRACE + 2) == [] # backing off + + state = event_apply( + DeploymentDeleted(deployment_id=deployment.deployment_id), state + ) + assert keeper.step(state, AFTER_GRACE + 3) == [] + assert keeper._attempts == {} # pyright: ignore[reportPrivateUsage] + state, _ = _deploy(state, MODEL_A) + + assert keeper.step(state, AFTER_GRACE + 4) != [] # a new deployment starts afresh + + +def test_deployments_created_and_deleted_leave_nothing_behind(): + keeper = Keeper(started_at=START) + state = _cluster(2) + now = AFTER_GRACE + for _ in range(100): + state, (deployment,) = _deploy(state, MODEL_A) + state = _apply(state, keeper.step(state, now)) + events, _ = keeper.delete( + DeleteDeployment(deployment_id=deployment.deployment_id), state + ) + state = _apply(state, events) + now += 1 + assert keeper.step(state, now) == [] + + assert state.instances == {} + assert state.deployments == {} + assert keeper._attempts == {} # pyright: ignore[reportPrivateUsage] + assert keeper._deleted == set() # pyright: ignore[reportPrivateUsage] + + +def test_delete_deployment_deletes_the_instance_the_keeper_placed(): + keeper = Keeper(started_at=START) + state, (deployment, _) = _deploy(_cluster(2), MODEL_A, MODEL_B) + state = _apply(state, keeper.step(state, AFTER_GRACE)) + state = _apply(state, keeper.step(state, AFTER_GRACE + 1)) + (kept,) = _instances_of(state, deployment.model_card) + (other,) = _instances_of(state, MODEL_B) + + events, downloads = keeper.delete( + DeleteDeployment(deployment_id=deployment.deployment_id), state + ) + state = _apply(state, events) + + assert deployment.deployment_id not in state.deployments + assert kept not in state.instances + assert other in state.instances + assert downloads == [] + assert keeper.step(state, AFTER_GRACE + 2) == [] + + +def test_delete_deployment_cancels_downloads_for_the_deleted_instance(): + keeper = Keeper(started_at=START) + state, (deployment,) = _deploy(_cluster(1), MODEL_A) + instance_id = _placed(events := keeper.step(state, AFTER_GRACE)) + state = _apply(state, events) + instance = state.instances[instance_id] + ((node_id, runner_id),) = instance.shard_assignments.node_to_runner.items() + downloading = DownloadOngoing( + node_id=node_id, + shard_metadata=instance.shard_assignments.runner_to_shard[runner_id], + download_progress=DownloadProgressData( + total=Memory.from_bytes(1000), + downloaded=Memory.from_bytes(300), + downloaded_this_session=Memory.from_bytes(300), + completed_files=0, + total_files=1, + speed=0.0, + eta_ms=0, + files={}, + ), + ) + state = state.model_copy(update={"downloads": {node_id: [downloading]}}) + + _, downloads = keeper.delete( + DeleteDeployment(deployment_id=deployment.deployment_id), state + ) + + assert downloads == [ + CancelDownload( + target_node_id=node_id, + model_id=MODEL_A.model_id, + command_id=downloads[0].command_id, + ) + ] + + +def test_delete_deployment_leaves_an_instance_someone_else_placed(): + keeper = Keeper(started_at=START) + state, (deployment,) = _deploy(_cluster(2), MODEL_A) + state, by_hand = _place_by_hand(state, MODEL_A) + assert keeper.step(state, AFTER_GRACE) == [] + + events, downloads = keeper.delete( + DeleteDeployment(deployment_id=deployment.deployment_id), state + ) + + assert _without_ids(events) == _without_ids( + [DeploymentDeleted(deployment_id=deployment.deployment_id)] + ) + assert downloads == [] + assert by_hand in _apply(state, events).instances + + +def test_delete_deployment_after_its_instance_is_gone(): + keeper = Keeper(started_at=START) + state, (deployment,) = _deploy(_cluster(2), MODEL_A) + instance_id = _placed(events := keeper.step(state, AFTER_GRACE)) + state = _lose(_apply(state, events), instance_id) + + events, downloads = keeper.delete( + DeleteDeployment(deployment_id=deployment.deployment_id), state + ) + + assert _without_ids(events) == _without_ids( + [DeploymentDeleted(deployment_id=deployment.deployment_id)] + ) + assert downloads == [] + + +def test_delete_deployment_deletes_an_instance_still_on_its_way(): + keeper = Keeper(started_at=START) + state, (deployment,) = _deploy(_cluster(2), MODEL_A) + placing = keeper.step(state, AFTER_GRACE) # sent, but not in the state yet + + deleting, _ = keeper.delete( + DeleteDeployment(deployment_id=deployment.deployment_id), state + ) + # The cluster applies the master's events in the order it sent them + state = _apply(state, [*placing, *deleting]) + + assert state.instances == {} + assert state.deployments == {} + + +def test_places_nothing_for_a_deployment_being_deleted(): + keeper = Keeper(started_at=START) + state, (deployment,) = _deploy(_cluster(2), MODEL_A) + + deleting, _ = keeper.delete( + DeleteDeployment(deployment_id=deployment.deployment_id), state + ) + # Until the deletion reaches the state, the state still shows the deployment + for now in (AFTER_GRACE, AFTER_GRACE + PLACEMENT_TIMEOUT, AFTER_GRACE + 3600): + assert keeper.step(state, now) == [] + state = _apply(state, deleting) + + assert keeper.step(state, AFTER_GRACE + 3601) == [] + assert keeper._deleted == set() # pyright: ignore[reportPrivateUsage] + + +def test_deleting_a_deployment_twice_is_harmless(): + keeper = Keeper(started_at=START) + state, (deployment,) = _deploy(_cluster(2), MODEL_A) + placing = keeper.step(state, AFTER_GRACE) + command = DeleteDeployment(deployment_id=deployment.deployment_id) + + first, _ = keeper.delete(command, state) + second, _ = keeper.delete(command, state) + state = _apply(state, [*placing, *first, *second]) + + assert state.instances == {} + assert state.deployments == {} + + +def test_delete_unknown_deployment_does_nothing(): + keeper = Keeper(started_at=START) + state, _ = _deploy(_cluster(1), MODEL_A) + + assert keeper.delete(DeleteDeployment(deployment_id=DeploymentId()), state) == ( + [], + [], + ) + + +# Status + + +def test_deployment_status_follows_its_instance(): + keeper = Keeper(started_at=START) + state, (deployment,) = _deploy(_cluster(2), MODEL_A) + + def status() -> str: + return deployment_status(state.deployments[deployment.deployment_id], state) + + assert status() == "placing" + instance_id = _placed(events := keeper.step(state, AFTER_GRACE)) + state = _apply(state, events) + assert status() == "starting" + state = _make_ready(state, instance_id) + assert status() == "serving" + (runner_id,) = state.instances[ + instance_id + ].shard_assignments.node_to_runner.values() + state = event_apply( + RunnerStatusUpdated(runner_id=runner_id, runner_status=RunnerRunning()), state + ) + assert status() == "serving" + state = event_apply( + RunnerStatusUpdated( + runner_id=runner_id, + runner_status=RunnerFailed(error_message="boom", diagnostics=[]), + ), + state, + ) + assert status() == "starting" + state = _lose(state, instance_id) + assert status() == "placing" + + +# Under random faults + + +def test_keeps_every_model_running_through_random_faults(): + """Lose instances and nodes at random for simulated hours. Throughout, the keeper only adds + instances, never two of a model at once; once the faults stop, every model is running again + within one backoff.""" + rng = random.Random(1234) + keeper = Keeper(started_at=START) + models = [_model(f"test/model-{i}") for i in range(4)] + full = _cluster(5) + state, _ = _deploy(full, *models) + now = AFTER_GRACE + + def step() -> None: + nonlocal state + events = keeper.step(state, now) + assert {type(e) for e in events} <= { + InstanceCreated, + DeploymentPlaced, + DeploymentPlacementFailed, + } + assert len([e for e in events if isinstance(e, InstanceCreated)]) <= 1 + state = _apply(state, events) + for model in models: + assert len(_instances_of(state, model)) <= 1 + + for _ in range(4 * 3600): + step() + fault = rng.random() + if fault < 0.01 and state.instances: + state = _lose(state, rng.choice(list(state.instances))) + elif fault < 0.02 and state.instances: + state = _make_ready(state, rng.choice(list(state.instances))) + elif fault < 0.022 and len(list(state.topology.list_nodes())) > 1: + dead = rng.choice(list(state.topology.list_nodes())) + on_dead = [ + instance_id + for instance_id, instance in state.instances.items() + if dead in instance.shard_assignments.node_to_runner + ] + state = _apply( + state, + [ + NodeTimedOut(node_id=dead), + *(InstanceDeleted(instance_id=i) for i in on_dead), + ], + ) + elif fault < 0.025: + # A node comes back: the full cluster again + state = state.model_copy( + update={ + "topology": full.topology, + "node_memory": full.node_memory, + "node_network": full.node_network, + "node_backends": full.node_backends, + } + ) + now += 1 + + for _ in range(int(MAX_BACKOFF + PLACEMENT_TIMEOUT) + len(models)): + step() + now += 1 + for model in models: + assert len(_instances_of(state, model)) == 1 + + +@pytest.mark.parametrize("seed", range(20)) +def test_deployments_come_and_go_while_the_state_trails_the_events(seed: int): + """As in the master, the keeper's events, deletions, creations and losses reach the state + only after a delay (here up to 10 s). Deployments are created and deleted, instances lost or + made ready, at random. At no point does a model have two instances, and once things settle, + exactly the models with a deployment have an instance.""" + rng = random.Random(seed) + keeper = Keeper(started_at=START) + models = [_model(f"test/model-{i}") for i in range(3)] + state = _cluster(3) + sent: list[tuple[float, Event]] = [] # sent by the master, not yet in the state + now = START + + def send(events: Iterable[Event]) -> None: + sent.extend((now, event) for event in events) + + def deliver(everything: bool) -> None: + nonlocal state + count = len(sent) if everything else rng.randint(0, 3) + # Nothing waits longer than 10 s + while sent and sent[0][0] < now - 10: + state = event_apply(sent.pop(0)[1], state) + for _ in range(min(count, len(sent))): + state = event_apply(sent.pop(0)[1], state) + for model in models: + assert len(_instances_of(state, model)) <= 1 + + for _ in range(2 * 3600): + action = rng.random() + if action < 0.02: + missing = [ + model + for model in models + if not any( + d.model_card.model_id == model.model_id + for d in state.deployments.values() + ) + ] + if missing: + send([DeploymentCreated(deployment=_deployment(rng.choice(missing)))]) + elif action < 0.03 and state.deployments: + deployment_id = rng.choice(list(state.deployments)) + events, _ = keeper.delete( + DeleteDeployment(deployment_id=deployment_id), state + ) + send(events) + elif action < 0.04 and state.instances: + send([InstanceDeleted(instance_id=rng.choice(list(state.instances)))]) + elif action < 0.06 and state.instances: + instance_id = rng.choice(list(state.instances)) + runner_ids = state.instances[ + instance_id + ].shard_assignments.node_to_runner.values() + send( + RunnerStatusUpdated(runner_id=runner_id, runner_status=RunnerReady()) + for runner_id in runner_ids + ) + send(keeper.step(state, now)) + deliver(everything=False) + now += 1 + + for _ in range(int(MAX_BACKOFF) + 10): + send(keeper.step(state, now)) + deliver(everything=True) + now += 1 + deployed = {d.model_card.model_id for d in state.deployments.values()} + for model in models: + expected = 1 if model.model_id in deployed else 0 + assert len(_instances_of(state, model)) == expected, model.model_id diff --git a/src/exo/master/tests/test_master_deployments.py b/src/exo/master/tests/test_master_deployments.py new file mode 100644 index 0000000000..44fae0b89b --- /dev/null +++ b/src/exo/master/tests/test_master_deployments.py @@ -0,0 +1,180 @@ +"""A real master keeping a deployment's model running, with its events going round through a +stand-in for the event router as they do in a node.""" + +from collections.abc import Callable +from datetime import datetime, timezone + +import anyio +import pytest + +from exo.master import keeper +from exo.master.main import Master +from exo.routing.router import get_node_zid +from exo.shared.models.model_cards import ModelCard, ModelTask +from exo.shared.types.backends import Backend +from exo.shared.types.commands import ( + Command, + CommandId, + CreateDeployment, + DeleteDeployment, + DeleteInstance, + ForwarderCommand, + ForwarderDownloadCommand, +) +from exo.shared.types.common import ModelId, SessionId, SystemId +from exo.shared.types.deployments import Deployment, DeploymentId +from exo.shared.types.events import ( + Event, + GlobalForwarderEvent, + LocalForwarderEvent, + NodeGatheredInfo, + RunnerStatusUpdated, +) +from exo.shared.types.memory import Memory +from exo.shared.types.profiling import MemoryUsage +from exo.shared.types.worker.instances import InstanceMeta +from exo.shared.types.worker.runners import RunnerReady +from exo.shared.types.worker.shards import Sharding +from exo.utils.channels import channel +from exo.utils.info_gatherer.info_gatherer import NodeBackends + +MODEL = ModelCard( + model_id=ModelId("test/model-a"), + n_layers=16, + storage_size=Memory.from_bytes(678948), + hidden_size=7168, + supports_tensor=True, + tasks=[ModelTask.TextGeneration], + backends=[Backend.MlxMetal], +) + + +async def _until(condition: Callable[[], bool]) -> None: + with anyio.fail_after(10): + while not condition(): + await anyio.sleep(0.01) + + +@pytest.mark.asyncio +async def test_master_keeps_a_deployed_model_running(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(keeper, "STARTUP_GRACE", 0.0) + node_id = get_node_zid() + session_id = SessionId(master_node_id=node_id, election_clock=0) + global_sender, _global_receiver = channel[GlobalForwarderEvent]() + command_sender, command_receiver = channel[ForwarderCommand]() + local_sender, local_receiver = channel[LocalForwarderEvent]() + download_sender, _download_receiver = channel[ForwarderDownloadCommand]() + event_sender, event_receiver = channel[Event]() + + async def event_router() -> None: + # The master's own events come back to it as local events, as through the real router + origin = SystemId() + with event_receiver as events: + index = 0 + async for event in events: + await local_sender.send( + LocalForwarderEvent( + origin=origin, origin_idx=index, session=session_id, event=event + ) + ) + index += 1 + + worker_index = 0 + + async def worker_reports(event: Event) -> None: + nonlocal worker_index + await local_sender.send( + LocalForwarderEvent( + origin=SystemId("Worker"), + origin_idx=worker_index, + session=session_id, + event=event, + ) + ) + worker_index += 1 + + async def node_info(info: MemoryUsage | NodeBackends) -> None: + await worker_reports( + NodeGatheredInfo( + when=str(datetime.now(tz=timezone.utc)), node_id=node_id, info=info + ) + ) + + async def send(command: Command) -> None: + await command_sender.send( + ForwarderCommand(origin=SystemId("API"), command=command) + ) + + master = Master( + node_id, + session_id, + event_sender=event_sender, + global_event_sender=global_sender, + local_event_receiver=local_receiver, + command_receiver=command_receiver, + download_command_sender=download_sender, + ) + async with anyio.create_task_group() as tg: + tg.start_soon(master.run) + tg.start_soon(event_router) + memory = Memory.from_bytes(678948 * 1024) + await node_info( + MemoryUsage( + ram_total=memory, + ram_available=memory, + swap_total=Memory.from_bytes(0), + swap_available=Memory.from_bytes(0), + ) + ) + await node_info(NodeBackends(backends=[Backend.MlxMetal])) + await _until(lambda: bool(master.state.node_backends)) + + deployment = Deployment( + deployment_id=DeploymentId(), + model_card=MODEL, + sharding=Sharding.Pipeline, + instance_meta=InstanceMeta.MlxRing, + min_nodes=1, + ) + await send(CreateDeployment(command_id=CommandId(), deployment=deployment)) + + # The keeper places an instance, and the deployment names it + await _until(lambda: bool(master.state.instances)) + (first,) = master.state.instances + await _until( + lambda: master.state.deployments[deployment.deployment_id].instance_id + == first, + ) + + # Its runner loads, and the keeper sees it ready + (runner_id,) = master.state.instances[ + first + ].shard_assignments.node_to_runner.values() + await worker_reports( + RunnerStatusUpdated(runner_id=runner_id, runner_status=RunnerReady()) + ) + await _until(lambda: runner_id in master.state.runners) + await anyio.sleep(1.5) + + # Someone deletes the instance: the keeper places another at once + await send(DeleteInstance(command_id=CommandId(), instance_id=first)) + await _until( + lambda: bool(master.state.instances) + and first not in master.state.instances, + ) + assert len(master.state.instances) == 1 + + # Deleting the deployment deletes its instance, and nothing is placed again + await send( + DeleteDeployment( + command_id=CommandId(), deployment_id=deployment.deployment_id + ) + ) + await _until( + lambda: not master.state.deployments and not master.state.instances, + ) + await anyio.sleep(2.5) + assert master.state.instances == {} + + event_sender.close() + await master.shutdown() diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py index 5f72247bac..caea161285 100644 --- a/src/exo/shared/apply.py +++ b/src/exo/shared/apply.py @@ -6,10 +6,15 @@ from exo.shared.models.model_cards import ModelCard from exo.shared.types.common import ModelId, NodeId +from exo.shared.types.deployments import Deployment, DeploymentId from exo.shared.types.events import ( ChunkGenerated, CustomModelCardAdded, CustomModelCardDeleted, + DeploymentCreated, + DeploymentDeleted, + DeploymentPlaced, + DeploymentPlacementFailed, Event, IndexedEvent, InputChunkReceived, @@ -95,6 +100,14 @@ def event_apply(event: Event, state: State) -> State: return apply_custom_model_card_added(event, state) case CustomModelCardDeleted(): return apply_custom_model_card_deleted(event, state) + case DeploymentCreated(): + return apply_deployment_created(event, state) + case DeploymentDeleted(): + return apply_deployment_deleted(event, state) + case DeploymentPlaced(): + return apply_deployment_placed(event, state) + case DeploymentPlacementFailed(): + return apply_deployment_placement_failed(event, state) case InstanceCreated(): return apply_instance_created(event, state) case InstanceDeleted(): @@ -500,3 +513,49 @@ def apply_custom_model_card_deleted( if model_id != event.model_id } return state.model_copy(update={"custom_model_cards": new_cards}) + + +def apply_deployment_created(event: DeploymentCreated, state: State) -> State: + # One deployment per model. Two requests for a model can both pass the API's check before + # either is applied; the first one applied wins. + model_id = event.deployment.model_card.model_id + if any(d.model_card.model_id == model_id for d in state.deployments.values()): + return state + deployments: Mapping[DeploymentId, Deployment] = { + **state.deployments, + event.deployment.deployment_id: event.deployment, + } + return state.model_copy(update={"deployments": deployments}) + + +def apply_deployment_deleted(event: DeploymentDeleted, state: State) -> State: + deployments: Mapping[DeploymentId, Deployment] = { + deployment_id: deployment + for deployment_id, deployment in state.deployments.items() + if deployment_id != event.deployment_id + } + return state.model_copy(update={"deployments": deployments}) + + +def apply_deployment_placed(event: DeploymentPlaced, state: State) -> State: + deployment = state.deployments.get(event.deployment_id) + if deployment is None: + return state + updated = deployment.model_copy( + update={"instance_id": event.instance_id, "placement_error": None} + ) + return state.model_copy( + update={"deployments": {**state.deployments, event.deployment_id: updated}} + ) + + +def apply_deployment_placement_failed( + event: DeploymentPlacementFailed, state: State +) -> State: + deployment = state.deployments.get(event.deployment_id) + if deployment is None: + return state + updated = deployment.model_copy(update={"placement_error": event.error}) + return state.model_copy( + update={"deployments": {**state.deployments, event.deployment_id: updated}} + ) diff --git a/src/exo/shared/tests/test_apply/test_apply_deployments.py b/src/exo/shared/tests/test_apply/test_apply_deployments.py new file mode 100644 index 0000000000..ebb9183e2f --- /dev/null +++ b/src/exo/shared/tests/test_apply/test_apply_deployments.py @@ -0,0 +1,133 @@ +from exo.shared.apply import event_apply +from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask +from exo.shared.types.backends import Backend +from exo.shared.types.deployments import Deployment, DeploymentId +from exo.shared.types.events import ( + DeploymentCreated, + DeploymentDeleted, + DeploymentPlaced, + DeploymentPlacementFailed, + IndexedEvent, +) +from exo.shared.types.memory import Memory +from exo.shared.types.state import State +from exo.shared.types.worker.instances import InstanceId, InstanceMeta +from exo.shared.types.worker.shards import Sharding + + +def _deployment(model_id: str) -> Deployment: + return Deployment( + deployment_id=DeploymentId(), + model_card=ModelCard( + model_id=ModelId(model_id), + storage_size=Memory.from_bytes(1000), + n_layers=10, + hidden_size=30, + supports_tensor=True, + tasks=[ModelTask.TextGeneration], + backends=[Backend.MlxMetal], + ), + sharding=Sharding.Pipeline, + instance_meta=InstanceMeta.MlxRing, + min_nodes=1, + ) + + +def test_created_deployment_is_added(): + deployment = _deployment("test/model-a") + + state = event_apply(DeploymentCreated(deployment=deployment), State()) + + assert state.deployments == {deployment.deployment_id: deployment} + + +def test_a_second_deployment_of_a_model_is_ignored(): + """Two requests can both pass the API's check before either is applied: the first wins.""" + first = _deployment("test/model-a") + state = event_apply(DeploymentCreated(deployment=first), State()) + + state = event_apply( + DeploymentCreated(deployment=_deployment("test/model-a")), state + ) + other = _deployment("test/model-b") + state = event_apply(DeploymentCreated(deployment=other), state) + + assert state.deployments == { + first.deployment_id: first, + other.deployment_id: other, + } + + +def test_deleted_deployment_is_removed(): + kept, deleted = _deployment("test/model-a"), _deployment("test/model-b") + state = State( + deployments={kept.deployment_id: kept, deleted.deployment_id: deleted} + ) + + state = event_apply(DeploymentDeleted(deployment_id=deleted.deployment_id), state) + state = event_apply(DeploymentDeleted(deployment_id=DeploymentId()), state) + + assert state.deployments == {kept.deployment_id: kept} + + +def test_placed_records_the_instance_and_clears_the_error(): + deployment = _deployment("test/model-a").model_copy( + update={"placement_error": "No cycles found with sufficient memory"} + ) + state = State(deployments={deployment.deployment_id: deployment}) + instance_id = InstanceId() + + state = event_apply( + DeploymentPlaced( + deployment_id=deployment.deployment_id, instance_id=instance_id + ), + state, + ) + + placed = state.deployments[deployment.deployment_id] + assert placed.instance_id == instance_id + assert placed.placement_error is None + + +def test_placement_failed_records_why_and_keeps_the_last_instance(): + instance_id = InstanceId() + deployment = _deployment("test/model-a").model_copy( + update={"instance_id": instance_id} + ) + state = State(deployments={deployment.deployment_id: deployment}) + + state = event_apply( + DeploymentPlacementFailed( + deployment_id=deployment.deployment_id, error="no room" + ), + state, + ) + + failed = state.deployments[deployment.deployment_id] + assert failed.placement_error == "no room" + assert failed.instance_id == instance_id + + +def test_events_for_a_deleted_deployment_change_nothing(): + """A placement can still be on its way when its deployment is deleted.""" + state = State() + + for event in ( + DeploymentPlaced(deployment_id=DeploymentId(), instance_id=InstanceId()), + DeploymentPlacementFailed(deployment_id=DeploymentId(), error="no room"), + ): + assert event_apply(event, state) == state + + +def test_deployments_survive_serialization(): + """The state is sent to nodes that join, and the events are stored in the event log.""" + deployment = _deployment("test/model-a").model_copy( + update={"instance_id": InstanceId(), "placement_error": "no room"} + ) + state = State(deployments={deployment.deployment_id: deployment}) + + restored = State.model_validate_json(state.model_dump_json()) + assert restored.deployments == state.deployments + + event = IndexedEvent(idx=0, event=DeploymentCreated(deployment=deployment)) + assert IndexedEvent.model_validate_json(event.model_dump_json()) == event diff --git a/src/exo/shared/types/commands.py b/src/exo/shared/types/commands.py index 67d318b255..7e85f6ea11 100644 --- a/src/exo/shared/types/commands.py +++ b/src/exo/shared/types/commands.py @@ -7,6 +7,7 @@ from exo.shared.models.model_cards import ModelCard, ModelId from exo.shared.types.chunks import InputImageChunk from exo.shared.types.common import CommandId, NodeId, SystemId +from exo.shared.types.deployments import Deployment, DeploymentId from exo.shared.types.instance_link import InstanceLinkId from exo.shared.types.text_generation import TextGenerationTaskParams from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta @@ -100,6 +101,14 @@ class DeleteInstanceLink(BaseCommand): link_id: InstanceLinkId +class CreateDeployment(BaseCommand): + deployment: Deployment + + +class DeleteDeployment(BaseCommand): + deployment_id: DeploymentId + + DownloadCommand = StartDownload | DeleteDownload | CancelDownload @@ -119,6 +128,8 @@ class DeleteInstanceLink(BaseCommand): | DeleteCustomModelCard | SetInstanceLink | DeleteInstanceLink + | CreateDeployment + | DeleteDeployment ) diff --git a/src/exo/shared/types/deployments.py b/src/exo/shared/types/deployments.py new file mode 100644 index 0000000000..3ad06b9b56 --- /dev/null +++ b/src/exo/shared/types/deployments.py @@ -0,0 +1,35 @@ +from typing import Literal + +from exo.shared.models.model_cards import ModelCard +from exo.shared.types.common import Id +from exo.shared.types.worker.instances import InstanceId, InstanceMeta +from exo.shared.types.worker.shards import Sharding +from exo.utils.pydantic_ext import FrozenModel + + +class DeploymentId(Id): + pass + + +class Deployment(FrozenModel): + """A request to keep a model running. + + While a deployment exists, the master's keeper places an instance of its model whenever the + cluster has none, with the same options a `/place_instance` request takes. Instances know + nothing about deployments: any instance of the model counts, however it was created. + """ + + deployment_id: DeploymentId + model_card: ModelCard + sharding: Sharding + instance_meta: InstanceMeta + min_nodes: int + # The instance the keeper placed last; deleting the deployment deletes it + instance_id: InstanceId | None = None + # Why the keeper's last placement failed, until a placement succeeds + placement_error: str | None = None + + +# serving: an instance of the model has every runner ready; starting: one is on its way up; +# placing: the keeper will place one shortly; cant_place: no placement fits the cluster right now +DeploymentStatus = Literal["serving", "starting", "placing", "cant_place"] diff --git a/src/exo/shared/types/events.py b/src/exo/shared/types/events.py index 01aa0ce5dc..35db6fb0cf 100644 --- a/src/exo/shared/types/events.py +++ b/src/exo/shared/types/events.py @@ -7,6 +7,7 @@ from exo.shared.topology import Connection from exo.shared.types.chunks import Chunk, InputImageChunk from exo.shared.types.common import CommandId, Id, ModelId, NodeId, SessionId, SystemId +from exo.shared.types.deployments import Deployment, DeploymentId from exo.shared.types.instance_link import InstanceLink, InstanceLinkId from exo.shared.types.tasks import Task, TaskId, TaskStatus from exo.shared.types.worker.downloads import DownloadProgress @@ -146,6 +147,28 @@ class InstanceLinkDeleted(BaseEvent): link_id: InstanceLinkId +class DeploymentCreated(BaseEvent): + deployment: Deployment + + +class DeploymentDeleted(BaseEvent): + deployment_id: DeploymentId + + +class DeploymentPlaced(BaseEvent): + """The keeper placed an instance for a deployment.""" + + deployment_id: DeploymentId + instance_id: InstanceId + + +class DeploymentPlacementFailed(BaseEvent): + """The keeper found no placement for a deployment's model; it tries again later.""" + + deployment_id: DeploymentId + error: str + + Event = ( TestEvent | TaskCreated @@ -169,6 +192,10 @@ class InstanceLinkDeleted(BaseEvent): | CustomModelCardDeleted | InstanceLinkCreated | InstanceLinkDeleted + | DeploymentCreated + | DeploymentDeleted + | DeploymentPlaced + | DeploymentPlacementFailed ) diff --git a/src/exo/shared/types/state.py b/src/exo/shared/types/state.py index baef1d72b9..c6097c1aac 100644 --- a/src/exo/shared/types/state.py +++ b/src/exo/shared/types/state.py @@ -9,6 +9,7 @@ from exo.shared.topology import Topology, TopologySnapshot from exo.shared.types.backends import Backend from exo.shared.types.common import ModelId, NodeId +from exo.shared.types.deployments import Deployment, DeploymentId from exo.shared.types.instance_link import InstanceLink, InstanceLinkId from exo.shared.types.profiling import ( DiskUsage, @@ -71,6 +72,9 @@ class State(FrozenModel): # User-added model cards. Workers can reconcile their on-disk custom card cache custom_model_cards: Mapping[ModelId, ModelCard] = {} + # Models someone asked to keep running; the master's keeper places an instance when one is missing + deployments: Mapping[DeploymentId, Deployment] = {} + @field_serializer("topology", mode="plain") def _encode_topology(self, value: Topology) -> TopologySnapshot: return value.to_snapshot()