mirror/exo
1
0
Fork 0
mirror of https://github.com/exo-explore/exo.git synced 2025-12-23 22:27:50 -05:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'master-node' into staging

Browse Source
This commit is contained in:
Arbion Halili
2025-07-15 16:32:08 +01:00
parent 21acd3794a 7fa7de8e83
commit d9b9aa7ad2
46 changed files with 1768 additions and 637 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
.vscode/extensions.json vendored Normal file
View File

@@ -0,0 +1,11 @@
{
"recommendations": [
"detachhead.basedpyright",
"ms-python.python"
],
"unwantedRecommendations": [
"ms-python.vscode-pylance",
"ms-python.pyright",
"ms-python.mypy-type-checker"
]
}

3
.vscode/settings.json vendored Normal file
View File

@@ -0,0 +1,3 @@
{
"basedpyright.importStrategy": "fromEnvironment"
}

29
.zed/settings.json Normal file
View File

@@ -0,0 +1,29 @@
// Folder-specific settings
//
// For a full list of overridable settings, and general information on folder-specific settings,
// see the documentation: https://zed.dev/docs/configuring-zed#settings-files
{
"lsp": {
"nix_python": {
"binary": {
"path": "nix",
"arguments": [
"run",
"--quiet",
"--no-warn-dirty",
"--no-allow-import-from-derivation",
"--print-build-logs",
"never",
"${projectRoot}#python-lsp",
"--",
"--stdio"
]
}
}
},
"languages": {
"Python": {
"language_servers": ["nix_python"]
}
}
}

14
flake.nix
View File

@@ -24,9 +24,23 @@
pkgs.protobuf
pkgs.rustc
pkgs.cargo
pkgs.basedpyright
pkgs.ruff
];
};
}
);
apps = forAllSystems (system:
let
pkgs = import nixpkgs { inherit system; };
in
{
python-lsp = {
type = "app";
program = "${pkgs.basedpyright}/bin/basedpyright-langserver";
};
}
);
};
}

4
justfile
View File

@@ -20,10 +20,10 @@ test:
uv run pytest master worker shared engines/*
check:
uv run basedpyright --project pyproject.toml
basedpyright --project pyproject.toml
sync:
uv sync --all-packages
uv sync --all-packages --reinstall
protobufs:
just regenerate-protobufs

21
master/commands.py Normal file
View File

@@ -0,0 +1,21 @@
from typing import Annotated, Literal
from pydantic import BaseModel, Field, TypeAdapter
class BaseExternalCommand[T: str](BaseModel):
command_type: T
class ChatCompletionNonStreamingCommand(
BaseExternalCommand[Literal["chat_completion_non_streaming"]]
):
command_type: Literal["chat_completion_non_streaming"] = (
"chat_completion_non_streaming"
)
ExternalCommand = Annotated[
ChatCompletionNonStreamingCommand, Field(discriminator="command_type")
]
ExternalCommandParser: TypeAdapter[ExternalCommand] = TypeAdapter(ExternalCommand)

5
master/env.py Normal file
View File

@@ -0,0 +1,5 @@
from shared.env import BaseEnv
class MasterEnvironmentSchema(BaseEnv):
pass

14
master/idempotency.py
View File

@@ -1,13 +1,13 @@
from hashlib import sha3_224 as hasher
from typing import Sequence, TypeVar
from typing import Sequence
from uuid import UUID
from shared.types.events.common import EventCategories, EventId, IdemKeyGenerator, State
EventCategoryT = TypeVar("EventCategoryT", bound=EventCategories)
from shared.types.events.common import EventCategory, EventId, IdemKeyGenerator, State
def get_idem_tag_generator(base: str) -> IdemKeyGenerator[EventCategoryT]:
def get_idem_tag_generator[EventCategoryT: EventCategory](
base: str,
) -> IdemKeyGenerator[EventCategoryT]:
"""Generates idempotency keys for events.
The keys are generated by hashing the state sequence number against a base string.
@@ -24,7 +24,9 @@ def get_idem_tag_generator(base: str) -> IdemKeyGenerator[EventCategoryT]:
*recurse(n - 1, next_hash),
)
initial_bytes = state.sequence_number.to_bytes(8, byteorder="big", signed=False)
initial_bytes = state.last_event_applied_idx.to_bytes(
8, byteorder="big", signed=False
)
return recurse(num_keys, initial_bytes)
return get_idem_keys

103
master/logging.py Normal file
View File

@@ -0,0 +1,103 @@
from collections.abc import Set
from typing import Literal
from shared.logging.common import LogEntry, LogEntryType
class MasterUninitializedLogEntry(LogEntry[Literal["master_uninitialized"]]):
entry_destination: Set[LogEntryType] = {LogEntryType.cluster}
entry_type: Literal["master_uninitialized"] = "master_uninitialized"
message: str = "No master state found, creating new one."
class MasterCommandReceivedLogEntry(LogEntry[Literal["master_command_received"]]):
entry_destination: Set[LogEntryType] = {LogEntryType.cluster}
entry_type: Literal["master_command_received"] = "master_command_received"
command_name: str
class MasterInvalidCommandReceivedLogEntry(
LogEntry[Literal["master_invalid_command_received"]]
):
entry_destination: Set[LogEntryType] = {LogEntryType.cluster}
entry_type: Literal["master_invalid_command_received"] = (
"master_invalid_command_received"
)
command_name: str
class MasterCommandRunnerNotRunningLogEntry: ...
class MasterStateManagerStoppedLogEntry: ...
class EventCategoryUnknownLogEntry(LogEntry[Literal["event_category_unknown"]]):
entry_destination: Set[LogEntryType] = {LogEntryType.cluster}
entry_type: Literal["event_category_unknown"] = "event_category_unknown"
event_category: str
message: str = "Event Category Unknown, Skipping Event."
class StateUpdateLoopAlreadyRunningLogEntry(
LogEntry[Literal["state_update_loop_already_running"]]
):
entry_destination: Set[LogEntryType] = {LogEntryType.cluster}
entry_type: Literal["state_update_loop_already_running"] = (
"state_update_loop_already_running"
)
message: str = "State Update Loop Already Running"
class StateUpdateLoopStartedLogEntry(LogEntry[Literal["state_update_loop_started"]]):
entry_destination: Set[LogEntryType] = {LogEntryType.cluster}
entry_type: Literal["state_update_loop_started"] = "state_update_loop_started"
message: str = "State Update Loop Started"
class StateUpdateLoopNotRunningLogEntry(
LogEntry[Literal["state_update_loop_not_running"]]
):
entry_destination: Set[LogEntryType] = {LogEntryType.cluster}
entry_type: Literal["state_update_loop_not_running"] = (
"state_update_loop_not_running"
)
message: str = "State Update Loop Not Running"
class StateUpdateLoopStoppedLogEntry(LogEntry[Literal["state_update_loop_stopped"]]):
entry_destination: Set[LogEntryType] = {LogEntryType.cluster}
entry_type: Literal["state_update_loop_stopped"] = "state_update_loop_stopped"
message: str = "State Update Loop Stopped"
class StateUpdateErrorLogEntry(LogEntry[Literal["state_update_error"]]):
entry_destination: Set[LogEntryType] = {LogEntryType.cluster}
entry_type: Literal["state_update_error"] = "state_update_error"
error: Exception
class StateUpdateEffectHandlerErrorLogEntry(
LogEntry[Literal["state_update_effect_handler_error"]]
):
entry_destination: Set[LogEntryType] = {LogEntryType.cluster}
entry_type: Literal["state_update_effect_handler_error"] = (
"state_update_effect_handler_error"
)
error: Exception
MasterLogEntries = (
MasterUninitializedLogEntry
| MasterCommandReceivedLogEntry
| MasterInvalidCommandReceivedLogEntry
| MasterCommandRunnerNotRunningLogEntry
| MasterStateManagerStoppedLogEntry
| EventCategoryUnknownLogEntry
| StateUpdateLoopAlreadyRunningLogEntry
| StateUpdateLoopStartedLogEntry
| StateUpdateLoopNotRunningLogEntry
| StateUpdateLoopStoppedLogEntry
| StateUpdateErrorLogEntry
| StateUpdateEffectHandlerErrorLogEntry
)

242
master/main.py
View File

@@ -1,6 +1,240 @@
def main():
print("Hello from master!")
from asyncio import CancelledError, Lock, Task, create_task
from asyncio import Queue as AsyncQueue
from contextlib import asynccontextmanager
from logging import Logger, LogRecord
from queue import Queue as PQueue
from typing import Callable, Sequence
from fastapi import FastAPI, Response
from fastapi.responses import StreamingResponse
from master.commands import ExternalCommand
from master.env import MasterEnvironmentSchema
from master.logging import (
MasterCommandRunnerNotRunningLogEntry,
MasterStateManagerStoppedLogEntry,
MasterUninitializedLogEntry,
)
from master.router import QueueMapping
from master.state_manager.sync import SyncStateManagerMapping
from shared.constants import EXO_MASTER_STATE
from shared.logger import (
FilterLogByType,
LogEntryType,
attach_to_queue,
configure_logger,
create_queue_listener,
log,
)
from shared.types.events.common import (
Apply,
EventCategory,
EventFromEventLog,
EventPublisher,
State,
)
from shared.types.models.common import ModelId
from shared.types.models.model import ModelInfo
from shared.types.states.master import MasterState
from shared.types.worker.common import InstanceId
from shared.types.worker.instances import Instance
if __name__ == "__main__":
main()
# Restore State
def get_master_state(logger: Logger) -> MasterState:
if EXO_MASTER_STATE.exists():
with open(EXO_MASTER_STATE, "r") as f:
return MasterState.model_validate_json(f.read())
else:
log(logger, MasterUninitializedLogEntry())
return MasterState()
# FastAPI Dependencies
def check_env_vars_defined(data: object, logger: Logger) -> MasterEnvironmentSchema:
if not isinstance(data, MasterEnvironmentSchema):
raise RuntimeError("Environment Variables Not Found")
return data
def get_master_state_dependency(data: object, logger: Logger) -> MasterState:
if not isinstance(data, MasterState):
raise RuntimeError("Master State Not Found")
return data
# Safety on Apply.
def safely_apply[T: EventCategory](
state: State[T], apply_fn: Apply[T], events: Sequence[EventFromEventLog[T]]
) -> State[T]:
sorted_events = sorted(events, key=lambda event: event.idx_in_log)
state = state.model_copy()
for event in sorted_events:
if event.idx_in_log <= state.last_event_applied_idx:
continue
state.last_event_applied_idx = event.idx_in_log
state = apply_fn(state, event)
return state
class MasterEventLoop:
"""Thread-safe manager for MasterState with independent event loop."""
def __init__(
self,
initial_state: MasterState,
push_events_to_queue: Callable[[QueueMapping], None],
event_publisher: EventPublisher[EventCategory],
state_managers: SyncStateManagerMapping,
logger: Logger,
):
self._state = initial_state
self._state_lock = Lock()
self._event_queues: QueueMapping
self._command_runner: ...
self._command_run_task: Task[None] | None = None
self._command_queue: AsyncQueue[ExternalCommand] = AsyncQueue()
self._response_queue: AsyncQueue[Response | StreamingResponse] = AsyncQueue()
self._state_managers: SyncStateManagerMapping
self._state_global_lock: Lock = Lock()
self._push_events_to_queue: Callable[[QueueMapping], None]
self._event_fetch_task: Task[None] | None = None
self._logger = logger
@property
def _is_command_runner_running(self) -> bool:
return self._command_run_task is not None and not self._command_run_task.done()
@property
def _is_event_fetcher_running(self) -> bool:
return self._event_fetch_task is not None and not self._event_fetch_task.done()
async def send_command(
self, command: ExternalCommand
) -> Response | StreamingResponse:
"""Send a command to the background event loop."""
if self._is_command_runner_running:
await self._command_queue.put(command)
return await self._response_queue.get()
else:
log(self._logger, MasterCommandRunnerNotRunningLogEntry())
raise RuntimeError("Command Runner Is Not Running")
async def start(self) -> None:
"""Start the background event loop."""
async def fetch_and_apply_events() -> None:
while True:
async with self._state_global_lock:
for state in self._state_managers.values():
self._push_events_to_queue(self._event_queues)
safely_apply(
state, apply_fn, self._event_queues[state.event_category]
)
self._event_fetch_task = create_task(fetch_and_apply_events())
self._command_run_task = create_task(self._command_runner())
async def stop(self) -> None:
"""Stop the background event loop and persist state."""
if not self._is_command_runner_running or not self._is_event_fetcher_running:
raise RuntimeError("Command Runner Is Not Running")
assert self._command_run_task is not None and self._event_fetch_task is not None
for service in [self._event_fetch_task, self._command_run_task]:
service.cancel()
try:
await service
except CancelledError:
pass
log(self._logger, MasterStateManagerStoppedLogEntry())
@asynccontextmanager
async def lifespan(app: FastAPI):
logger = configure_logger("master")
telemetry_queue: PQueue[LogRecord] = PQueue()
metrics_queue: PQueue[LogRecord] = PQueue()
cluster_queue: PQueue[LogRecord] = PQueue()
attach_to_queue(
logger,
[
FilterLogByType(log_types={LogEntryType.telemetry}),
],
telemetry_queue,
)
attach_to_queue(
logger,
[
FilterLogByType(log_types={LogEntryType.metrics}),
],
metrics_queue,
)
attach_to_queue(
logger,
[
FilterLogByType(log_types={LogEntryType.cluster}),
],
cluster_queue,
)
# TODO: Add handlers
telemetry_listener = create_queue_listener(telemetry_queue, [])
metrics_listener = create_queue_listener(metrics_queue, [])
cluster_listener = create_queue_listener(cluster_queue, [])
telemetry_listener.start()
metrics_listener.start()
cluster_listener.start()
initial_state = get_master_state(logger)
app.state.master_event_loop = MasterEventLoop(
initial_state, None, None, None, logger
)
await app.state.master_event_loop.start()
yield
await app.state.master_event_loop.stop()
app = FastAPI(lifespan=lifespan)
@app.get("/topology/control_plane")
def get_control_plane_topology():
return {"message": "Hello, World!"}
@app.get("/topology/data_plane")
def get_data_plane_topology():
return {"message": "Hello, World!"}
@app.get("/instances/list")
def list_instances():
return {"message": "Hello, World!"}
@app.post("/instances/create")
def create_instance(model_id: ModelId) -> InstanceId: ...
@app.get("/instance/{instance_id}/read")
def get_instance(instance_id: InstanceId) -> Instance: ...
@app.delete("/instance/{instance_id}/delete")
def remove_instance(instance_id: InstanceId) -> None: ...
@app.get("/model/{model_id}/metadata")
def get_model_data(model_id: ModelId) -> ModelInfo: ...
@app.post("/model/{model_id}/instances")
def get_instances_by_model(model_id: ModelId) -> list[Instance]: ...

7
master/pyproject.toml
View File

@@ -4,7 +4,10 @@ version = "0.1.0"
description = "Master service for the Exo project"
readme = "README.md"
requires-python = ">=3.13"
dependencies = ["exo-shared"]
dependencies = [
"exo-shared",
"fastapi>=0.116.0",
]
[build-system]
requires = ["hatchling"]
@@ -21,4 +24,4 @@ exclude = ["*.md", "pyproject.toml"]
[tool.hatch.build.targets.sdist]
packages = []
include = ["*"]
exclude = ["*.md", "pyproject.toml"]
exclude = ["*.md", "pyproject.toml"]

90
master/router.py Normal file
View File

@@ -0,0 +1,90 @@
from asyncio import Queue, gather
from logging import Logger
from typing import Literal, Protocol, TypedDict
from master.sanity_checking import check_keys_in_map_match_enum_values
from shared.types.events.common import (
EventCategories,
EventCategory,
EventCategoryEnum,
EventFromEventLog,
narrow_event_from_event_log_type,
)
class QueueMapping(TypedDict):
MutatesTaskState: Queue[
EventFromEventLog[Literal[EventCategoryEnum.MutatesTaskState]]
]
MutatesTaskSagaState: Queue[
EventFromEventLog[Literal[EventCategoryEnum.MutatesTaskSagaState]]
]
MutatesControlPlaneState: Queue[
EventFromEventLog[Literal[EventCategoryEnum.MutatesControlPlaneState]]
]
MutatesDataPlaneState: Queue[
EventFromEventLog[Literal[EventCategoryEnum.MutatesDataPlaneState]]
]
MutatesRunnerStatus: Queue[
EventFromEventLog[Literal[EventCategoryEnum.MutatesRunnerStatus]]
]
MutatesInstanceState: Queue[
EventFromEventLog[Literal[EventCategoryEnum.MutatesInstanceState]]
]
MutatesNodePerformanceState: Queue[
EventFromEventLog[Literal[EventCategoryEnum.MutatesNodePerformanceState]]
]
check_keys_in_map_match_enum_values(QueueMapping, EventCategoryEnum)
class EventRouterProtocol(Protocol):
queue_map: QueueMapping
start_idx: int
def sync_queues(self) -> None: ...
class EventRouter(EventRouterProtocol):
"""Routes events to appropriate services based on event categories."""
queue_map: QueueMapping
start_idx: int
logger: Logger
async def _get_queue_by_category[T: EventCategory](
self, category: T
) -> Queue[EventFromEventLog[T]]:
"""Get the queue for a given category."""
category_str: str = category.value
queue: Queue[EventFromEventLog[T]] = self.queue_map[category_str]
return queue
async def _process_events[T: EventCategory](self, category: T) -> None:
"""Process events for a given domain."""
queue: Queue[EventFromEventLog[T]] = await self._get_queue_by_category(category)
events_to_process: list[EventFromEventLog[T]] = []
while not queue.empty():
events_to_process.append(await queue.get())
for event_to_process in events_to_process:
await self.queue_map[category.value].put(event_to_process)
return None
async def _submit_events[T: EventCategory | EventCategories](
self, events: list[EventFromEventLog[T]]
) -> None:
"""Route multiple events to their appropriate services."""
for event in events:
if isinstance(event.event.event_category, EventCategory):
q1: Queue[EventFromEventLog[T]] = self.queue_map[
event.event.event_category.value
]
await q1.put(event)
elif isinstance(event.event.event_category, EventCategories):
for category in event.event.event_category:
narrow_event = narrow_event_from_event_log_type(event, category)
q2: Queue[EventFromEventLog[T]] = self.queue_map[category.value]
await q2.put(narrow_event)
await gather(*[self._process_events(domain) for domain in EventCategoryEnum])

13
master/sanity_checking.py Normal file
View File

@@ -0,0 +1,13 @@
from enum import StrEnum
from typing import Any, Mapping, Type
def check_keys_in_map_match_enum_values[TEnum: StrEnum](
mapping_type: Type[Mapping[Any, Any]],
enum: Type[TEnum],
) -> None:
mapping_keys = set(mapping_type.__annotations__.keys())
category_values = set(e.value for e in enum)
assert mapping_keys == category_values, (
f"StateDomainMapping keys {mapping_keys} do not match EventCategories values {category_values}"
)

128
master/state_manager/async.py Normal file
View File

@@ -0,0 +1,128 @@
from asyncio import Lock, Queue, Task, create_task
from logging import Logger
from typing import List, Literal, Protocol, TypedDict
from master.logging import (
StateUpdateEffectHandlerErrorLogEntry,
StateUpdateErrorLogEntry,
StateUpdateLoopAlreadyRunningLogEntry,
StateUpdateLoopNotRunningLogEntry,
StateUpdateLoopStartedLogEntry,
StateUpdateLoopStoppedLogEntry,
)
from master.router import check_keys_in_map_match_enum_values
from shared.constants import EXO_ERROR_REPORTING_MESSAGE
from shared.logger import log
from shared.types.events.common import (
Apply,
EffectHandler,
EventCategory,
EventCategoryEnum,
EventFromEventLog,
State,
StateAndEvent,
)
class AsyncStateManager[EventCategoryT: EventCategory](Protocol):
"""Protocol for services that manage a specific state domain."""
_task: Task[None] | None
_logger: Logger
_apply: Apply[EventCategoryT]
_default_effects: List[EffectHandler[EventCategoryT]]
extra_effects: List[EffectHandler[EventCategoryT]]
state: State[EventCategoryT]
queue: Queue[EventFromEventLog[EventCategoryT]]
lock: Lock
def __init__(
self,
state: State[EventCategoryT],
queue: Queue[EventFromEventLog[EventCategoryT]],
extra_effects: List[EffectHandler[EventCategoryT]],
logger: Logger,
) -> None:
"""Initialise the service with its event queue."""
self.state = state
self.queue = queue
self.extra_effects = extra_effects
self._logger = logger
self._task = None
async def read_state(self) -> State[EventCategoryT]:
"""Get a thread-safe snapshot of this service's state domain."""
return self.state.model_copy(deep=True)
@property
def is_running(self) -> bool:
"""Check if the service's event loop is running."""
return self._task is not None and not self._task.done()
async def start(self) -> None:
"""Start the service's event loop."""
if self.is_running:
log(self._logger, StateUpdateLoopAlreadyRunningLogEntry())
raise RuntimeError("State Update Loop Already Running")
log(self._logger, StateUpdateLoopStartedLogEntry())
self._task = create_task(self._event_loop())
async def stop(self) -> None:
"""Stop the service's event loop."""
if not self.is_running:
log(self._logger, StateUpdateLoopNotRunningLogEntry())
raise RuntimeError("State Update Loop Not Running")
assert self._task is not None, (
f"{EXO_ERROR_REPORTING_MESSAGE()}"
"BUG: is_running is True but _task is None, this should never happen!"
)
self._task.cancel()
log(self._logger, StateUpdateLoopStoppedLogEntry())
async def _event_loop(self) -> None:
"""Event loop for the service."""
while True:
event = await self.queue.get()
previous_state = self.state.model_copy(deep=True)
try:
async with self.lock:
updated_state = self._apply(
self.state,
event,
)
self.state = updated_state
except Exception as e:
log(self._logger, StateUpdateErrorLogEntry(error=e))
raise e
try:
for effect_handler in self._default_effects + self.extra_effects:
effect_handler(StateAndEvent(previous_state, event), updated_state)
except Exception as e:
log(self._logger, StateUpdateEffectHandlerErrorLogEntry(error=e))
raise e
class AsyncStateManagerMapping(TypedDict):
MutatesTaskState: AsyncStateManager[Literal[EventCategoryEnum.MutatesTaskState]]
MutatesTaskSagaState: AsyncStateManager[
Literal[EventCategoryEnum.MutatesTaskSagaState]
]
MutatesControlPlaneState: AsyncStateManager[
Literal[EventCategoryEnum.MutatesControlPlaneState]
]
MutatesDataPlaneState: AsyncStateManager[
Literal[EventCategoryEnum.MutatesDataPlaneState]
]
MutatesRunnerStatus: AsyncStateManager[
Literal[EventCategoryEnum.MutatesRunnerStatus]
]
MutatesInstanceState: AsyncStateManager[
Literal[EventCategoryEnum.MutatesInstanceState]
]
MutatesNodePerformanceState: AsyncStateManager[
Literal[EventCategoryEnum.MutatesNodePerformanceState]
]
check_keys_in_map_match_enum_values(AsyncStateManagerMapping, EventCategoryEnum)

19
master/state_manager/sync.py Normal file
View File

@@ -0,0 +1,19 @@
from typing import Literal, TypedDict
from master.sanity_checking import check_keys_in_map_match_enum_values
from shared.types.events.common import EventCategoryEnum, State
class SyncStateManagerMapping(TypedDict):
MutatesTaskState: State[Literal[EventCategoryEnum.MutatesTaskState]]
MutatesTaskSagaState: State[Literal[EventCategoryEnum.MutatesTaskSagaState]]
MutatesControlPlaneState: State[Literal[EventCategoryEnum.MutatesControlPlaneState]]
MutatesDataPlaneState: State[Literal[EventCategoryEnum.MutatesDataPlaneState]]
MutatesRunnerStatus: State[Literal[EventCategoryEnum.MutatesRunnerStatus]]
MutatesInstanceState: State[Literal[EventCategoryEnum.MutatesInstanceState]]
MutatesNodePerformanceState: State[
Literal[EventCategoryEnum.MutatesNodePerformanceState]
]
check_keys_in_map_match_enum_values(SyncStateManagerMapping, EventCategoryEnum)

1
networking/src/networking/_core.pyi
View File

@@ -1 +0,0 @@
def hello_from_bin() -> str: ...

1
pyproject.toml
View File

@@ -12,7 +12,6 @@ dependencies = [
# dependencies only required for development
[dependency-groups]
dev = [
"basedpyright>=1.29.4",
"maturin>=1.9.0",
"pytest>=8.4.0",
"pytest-asyncio>=1.0.0",

20
shared/constants.py
View File

@@ -1,11 +1,27 @@
import inspect
from pathlib import Path
EXO_HOME = Path.home() / ".exo"
EXO_EVENT_DB = EXO_HOME / "event_db.sqlite3"
EXO_MASTER_CONFIG = EXO_HOME / "master.json"
EXO_WORKER_CONFIG = EXO_HOME / "worker.json"
EXO_MASTER_STATE = EXO_HOME / "master_state.json"
EXO_WORKER_STATE = EXO_HOME / "worker_state.json"
EXO_MASTER_LOG = EXO_HOME / "master.log"
EXO_WORKER_LOG = EXO_HOME / "worker.log"
EXO_WORKER_KEYRING_FILE = EXO_HOME / "worker_keyring"
EXO_MASTER_KEYRING_FILE = EXO_HOME / "master_keyring"
# little helper function to get the name of the module that raised the error
def get_caller_module_name() -> str:
frm = inspect.stack()[1]
mod = inspect.getmodule(frm[0])
if mod is None:
return "UNKNOWN MODULE"
return mod.__name__
EXO_ERROR_REPORTING_MESSAGE = lambda: (
f"THIS IS A BUG IN THE EXO SOFTWARE, PLEASE REPORT IT AT https://github.com/exo-explore/exo/\n"
f"The module that raised the error was: {get_caller_module_name()}"
)

221
shared/graphs/networkx.py Normal file
View File

@@ -0,0 +1,221 @@
from dataclasses import dataclass
from typing import Mapping, Set
import rustworkx as rx
from pydantic import TypeAdapter
from shared.types.graphs.common import (
Edge,
EdgeData,
EdgeIdT,
EdgeTypeT,
MutableGraphProtocol,
Vertex,
VertexData,
VertexIdT,
VertexTypeT,
)
@dataclass(frozen=True)
class _VertexWrapper[VertexTypeT, VertexIdT]:
"""Internal wrapper to store vertex ID alongside vertex data."""
vertex_id: VertexIdT
vertex_data: VertexData[VertexTypeT]
@dataclass(frozen=True)
class _EdgeWrapper[EdgeTypeT, EdgeIdT]:
"""Internal wrapper to store edge ID alongside edge data."""
edge_id: EdgeIdT
edge_data: EdgeData[EdgeTypeT]
class NetworkXGraph(MutableGraphProtocol[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]):
edge_base: TypeAdapter[EdgeTypeT]
vertex_base: TypeAdapter[VertexTypeT]
_graph: rx.PyDiGraph[
_VertexWrapper[VertexTypeT, VertexIdT], _EdgeWrapper[EdgeTypeT, EdgeIdT]
]
_vertex_id_to_index: dict[VertexIdT, int]
_edge_id_to_endpoints: dict[EdgeIdT, tuple[int, int]]
def __init__(
self, edge_base: TypeAdapter[EdgeTypeT], vertex_base: TypeAdapter[VertexTypeT]
) -> None:
self.edge_base = edge_base
self.vertex_base = vertex_base
self._graph = rx.PyDiGraph()
self._vertex_id_to_index = {}
self._edge_id_to_endpoints = {}
###
# GraphProtocol methods
###
def list_edges(self) -> Set[EdgeIdT]:
return {edge.edge_id for edge in self._graph.edges()}
def list_vertices(self) -> Set[VertexIdT]:
return {node.vertex_id for node in self._graph.nodes()}
def get_vertices_from_edges(
self, edges: Set[EdgeIdT]
) -> Mapping[EdgeIdT, Set[VertexIdT]]:
result: dict[EdgeIdT, Set[VertexIdT]] = {}
for edge_id in edges:
if edge_id in self._edge_id_to_endpoints:
u_idx, v_idx = self._edge_id_to_endpoints[edge_id]
u_data = self._graph.get_node_data(u_idx)
v_data = self._graph.get_node_data(v_idx)
result[edge_id] = {u_data.vertex_id, v_data.vertex_id}
return result
def get_edges_from_vertices(
self, vertices: Set[VertexIdT]
) -> Mapping[VertexIdT, Set[EdgeIdT]]:
result: dict[VertexIdT, Set[EdgeIdT]] = {}
for vertex_id in vertices:
if vertex_id in self._vertex_id_to_index:
vertex_idx = self._vertex_id_to_index[vertex_id]
edge_ids: Set[EdgeIdT] = set()
# Get outgoing edges
for _, _, edge_data in self._graph.out_edges(vertex_idx):
edge_ids.add(edge_data.edge_id)
# Get incoming edges
for _, _, edge_data in self._graph.in_edges(vertex_idx):
edge_ids.add(edge_data.edge_id)
result[vertex_id] = edge_ids
return result
def get_edge_data(
self, edges: Set[EdgeIdT]
) -> Mapping[EdgeIdT, EdgeData[EdgeTypeT]]:
result: dict[EdgeIdT, EdgeData[EdgeTypeT]] = {}
for edge_id in edges:
if edge_id in self._edge_id_to_endpoints:
u_idx, v_idx = self._edge_id_to_endpoints[edge_id]
edge_wrapper = self._graph.get_edge_data(u_idx, v_idx)
result[edge_id] = edge_wrapper.edge_data
return result
def get_vertex_data(
self, vertices: Set[VertexIdT]
) -> Mapping[VertexIdT, VertexData[VertexTypeT]]:
result: dict[VertexIdT, VertexData[VertexTypeT]] = {}
for vertex_id in vertices:
if vertex_id in self._vertex_id_to_index:
vertex_idx = self._vertex_id_to_index[vertex_id]
vertex_wrapper = self._graph.get_node_data(vertex_idx)
result[vertex_id] = vertex_wrapper.vertex_data
return result
###
# MutableGraphProtocol methods
###
def check_edges_exists(self, edge_id: EdgeIdT) -> bool:
return edge_id in self._edge_id_to_endpoints
def check_vertex_exists(self, vertex_id: VertexIdT) -> bool:
return vertex_id in self._vertex_id_to_index
def _add_edge(self, edge_id: EdgeIdT, edge_data: EdgeData[EdgeTypeT]) -> None:
# This internal method is not used in favor of a safer `attach_edge` implementation.
raise NotImplementedError(
"Use attach_edge to add edges. The internal _add_edge protocol method is flawed."
)
def _add_vertex(
self, vertex_id: VertexIdT, vertex_data: VertexData[VertexTypeT]
) -> None:
if vertex_id not in self._vertex_id_to_index:
wrapper = _VertexWrapper(vertex_id=vertex_id, vertex_data=vertex_data)
idx = self._graph.add_node(wrapper)
self._vertex_id_to_index[vertex_id] = idx
def _remove_edge(self, edge_id: EdgeIdT) -> None:
if edge_id in self._edge_id_to_endpoints:
u_idx, v_idx = self._edge_id_to_endpoints[edge_id]
self._graph.remove_edge(u_idx, v_idx)
del self._edge_id_to_endpoints[edge_id]
else:
raise ValueError(f"Edge with id {edge_id} not found.")
def _remove_vertex(self, vertex_id: VertexIdT) -> None:
if vertex_id in self._vertex_id_to_index:
vertex_idx = self._vertex_id_to_index[vertex_id]
# Remove any edges connected to this vertex from our mapping
edges_to_remove: list[EdgeIdT] = []
for edge_id, (u_idx, v_idx) in self._edge_id_to_endpoints.items():
if u_idx == vertex_idx or v_idx == vertex_idx:
edges_to_remove.append(edge_id)
for edge_id in edges_to_remove:
del self._edge_id_to_endpoints[edge_id]
# Remove the vertex from the graph
self._graph.remove_node(vertex_idx)
del self._vertex_id_to_index[vertex_id]
else:
raise ValueError(f"Vertex with id {vertex_id} not found.")
def attach_edge(
self,
edge: Edge[EdgeTypeT, EdgeIdT, VertexIdT],
extra_vertex: Vertex[VertexTypeT, EdgeIdT, VertexIdT] | None = None,
) -> None:
"""
Attaches an edge to the graph, overriding the default protocol implementation.
This implementation corrects a flaw in the protocol's `_add_edge`
signature and provides more intuitive behavior when connecting existing vertices.
"""
base_vertex_id, target_vertex_id = edge.edge_vertices
if not self.check_vertex_exists(base_vertex_id):
raise ValueError(f"Base vertex {base_vertex_id} does not exist.")
target_vertex_exists = self.check_vertex_exists(target_vertex_id)
if not target_vertex_exists:
if extra_vertex is None:
raise ValueError(
f"Target vertex {target_vertex_id} does not exist and no `extra_vertex` was provided."
)
if extra_vertex.vertex_id != target_vertex_id:
raise ValueError(
f"The ID of `extra_vertex` ({extra_vertex.vertex_id}) does not match "
f"the target vertex ID of the edge ({target_vertex_id})."
)
self._add_vertex(extra_vertex.vertex_id, extra_vertex.vertex_data)
elif extra_vertex is not None:
raise ValueError(
f"Target vertex {target_vertex_id} already exists, but `extra_vertex` was provided."
)
# Get the internal indices
base_idx = self._vertex_id_to_index[base_vertex_id]
target_idx = self._vertex_id_to_index[target_vertex_id]
# Create edge wrapper and add to graph
edge_wrapper = _EdgeWrapper(edge_id=edge.edge_id, edge_data=edge.edge_data)
self._graph.add_edge(base_idx, target_idx, edge_wrapper)
# Store the mapping
self._edge_id_to_endpoints[edge.edge_id] = (base_idx, target_idx)

30
shared/logger.py
View File

@@ -1,12 +1,32 @@
import logging
import logging.handlers
from collections.abc import Sequence, Set
from enum import Enum
from queue import Queue
from typing import Annotated
from pydantic import BaseModel
from pydantic import Field, TypeAdapter
from rich.logging import RichHandler
from master.logging import MasterLogEntries
from shared.logging.common import LogEntryType
from worker.logging import WorkerLogEntries
LogEntries = Annotated[
MasterLogEntries | WorkerLogEntries, Field(discriminator="entry_type")
]
LogParser: TypeAdapter[LogEntries] = TypeAdapter(LogEntries)
class FilterLogByType(logging.Filter):
def __init__(self, log_types: Set[LogEntryType]):
super().__init__()
self.log_types = log_types
def filter(self, record: logging.LogRecord) -> bool:
message = record.getMessage()
LogParser.validate_json(message)
return True
class LogEntryType(str, Enum):
telemetry = "telemetry"
@@ -79,3 +99,9 @@ def create_queue_listener(
log_queue, *effect_handlers, respect_handler_level=True
)
return listener
def log(
logger: logging.Logger, log_entry: LogEntries, log_level: int = logging.INFO
) -> None:
logger.log(log_level, log_entry.model_dump_json())

18
shared/logging/common.py Normal file
View File

@@ -0,0 +1,18 @@
from collections.abc import Set
from enum import Enum
from typing import Generic, TypeVar
from pydantic import BaseModel
LogEntryTypeT = TypeVar("LogEntryTypeT", bound=str)
class LogEntryType(str, Enum):
telemetry = "telemetry"
metrics = "metrics"
cluster = "cluster"
class LogEntry(BaseModel, Generic[LogEntryTypeT]):
entry_destination: Set[LogEntryType]
entry_type: LogEntryTypeT

18
shared/openai.py
View File

@@ -1,21 +1,23 @@
from typing import TYPE_CHECKING, Literal, TypeAlias, get_type_hints
FinishReason: TypeAlias = Literal[
"stop", "length", "tool_calls", "content_filter", "function_call"
]
if TYPE_CHECKING:
import openai.types as openai_types
import openai.types.chat as openai_chat
types = openai_types
chat = openai_chat
assert (
get_type_hints(chat.chat_completion_chunk.Choice)["finish_reason"] == FinishReason
), "Upstream changed Choice.finish_reason; update FinishReason alias."
else:
types = None
chat = None
FinishReason: TypeAlias = Literal[
"stop", "length", "tool_calls", "content_filter", "function_call"
]
if TYPE_CHECKING:
assert (
get_type_hints(chat.chat_completion_chunk.Choice)["finish_reason"]
== FinishReason
), "Upstream changed Choice.finish_reason; update FinishReason alias."
__all__ = ["types", "chat", "FinishReason"]

2
shared/pyproject.toml
View File

@@ -5,11 +5,13 @@ description = "Shared utilities for the Exo project"
readme = "README.md"
requires-python = ">=3.13"
dependencies = [
"networkx>=3.5",
"openai>=1.93.0",
"pathlib>=1.0.1",
"protobuf>=6.31.1",
"pydantic>=2.11.7",
"rich>=14.0.0",
"rustworkx>=0.16.0",
]
[build-system]

20
shared/types/events/chunks.py
View File

@@ -1,28 +1,22 @@
from enum import Enum
from typing import Annotated, Generic, Literal, TypeVar
from typing import Annotated, Literal
from openai.types.chat.chat_completion import ChatCompletion
from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
# from openai.types.chat.chat_completion import ChatCompletion
# from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
from pydantic import BaseModel, Field, TypeAdapter
from shared.openai import FinishReason
from shared.types.models.common import ModelId
from shared.types.tasks.common import TaskId
OpenAIResponse = (
ChatCompletion | ChatCompletionChunk
) ## Currently we only support chat completions
class ChunkType(str, Enum):
token = "token"
image = "image"
ChunkT = TypeVar("ChunkT", bound=ChunkType)
class BaseChunk(BaseModel, Generic[ChunkT]):
class BaseChunk[ChunkTypeT: ChunkType](BaseModel):
chunk_type: ChunkTypeT
task_id: TaskId
idx: int
model: ModelId
@@ -59,6 +53,10 @@ class ImageChunk(BaseChunk[ChunkType.image]):
GenerationChunk = Annotated[TokenChunk | ImageChunk, Field(discriminator="chunk_type")]
GenerationChunkTypeAdapter: TypeAdapter[GenerationChunk] = TypeAdapter(GenerationChunk)
## OpenAIResponse = (
## ChatCompletion | ChatCompletionChunk
## ) ## Currently we only support chat completions
# my_chunk: dict[str, Any] = TokenChunk(
# task_id=TaskId('nicerid'),
# idx=0,

296
shared/types/events/common.py
View File

@@ -1,17 +1,21 @@
from enum import Enum, auto
from enum import Enum, StrEnum
from typing import (
Annotated,
Callable,
Generic,
FrozenSet,
Literal,
NamedTuple,
Protocol,
Sequence,
Tuple,
TypeVar,
cast,
)
from pydantic import BaseModel, Field, TypeAdapter, model_validator
from pydantic import BaseModel, Field, model_validator
from shared.types.common import NewUUID, NodeId
from shared.types.events.sanity_checking import (
assert_literal_union_covers_enum,
check_event_type_union_is_consistent_with_registry,
)
class EventId(NewUUID):
@@ -22,14 +26,16 @@ class TimerId(NewUUID):
pass
class MLXEventTypes(str, Enum):
# Here are all the unique kinds of events that can be sent over the network.
# I've defined them in different enums for clarity, but they're all part of the same set of possible events.
class TaskSagaEventTypes(str, Enum):
MLXInferenceSagaPrepare = "MLXInferenceSagaPrepare"
MLXInferenceSagaStartPrepare = "MLXInferenceSagaStartPrepare"
class TaskEventTypes(str, Enum):
TaskCreated = "TaskCreated"
TaskUpdated = "TaskUpdated"
TaskStateUpdated = "TaskStateUpdated"
TaskDeleted = "TaskDeleted"
@@ -40,22 +46,22 @@ class StreamingEventTypes(str, Enum):
class InstanceEventTypes(str, Enum):
InstanceCreated = "InstanceCreated"
InstanceDeleted = "InstanceDeleted"
InstanceToBeReplacedAtomically = "InstanceToBeReplacedAtomically"
InstanceActivated = "InstanceActivated"
InstanceDeactivated = "InstanceDeactivated"
InstanceReplacedAtomically = "InstanceReplacedAtomically"
InstanceStatusUpdated = "InstanceStatusUpdated"
class InstanceStateEventTypes(str, Enum):
InstanceRunnerStateUpdated = "InstanceRunnerStateUpdated"
class RunnerStatusEventTypes(str, Enum):
RunnerStatusUpdated = "RunnerStatusUpdated"
class NodePerformanceEventTypes(str, Enum):
NodePerformanceProfiled = "NodePerformanceProfiled"
NodePerformanceMeasured = "NodePerformanceMeasured"
class DataPlaneEventTypes(str, Enum):
DataPlaneEdgeCreated = "DataPlaneEdgeCreated"
DataPlaneEdgeProfiled = "DataPlaneEdgeProfiled"
DataPlaneEdgeReplacedAtomically = "DataPlaneEdgeReplacedAtomically"
DataPlaneEdgeDeleted = "DataPlaneEdgeDeleted"
@@ -70,167 +76,154 @@ class TimerEventTypes(str, Enum):
TimerFired = "TimerFired"
class ResourceEventTypes(str, Enum):
ResourceProfiled = "ResourceProfiled"
# Registry of all event type enums
EVENT_TYPE_ENUMS = [
TaskEventTypes,
StreamingEventTypes,
InstanceEventTypes,
RunnerStatusEventTypes,
NodePerformanceEventTypes,
DataPlaneEventTypes,
ControlPlaneEventTypes,
TimerEventTypes,
TaskSagaEventTypes,
]
class EventCategories(str, Enum):
TaskEventTypes = auto()
StreamingEventTypes = auto()
InstanceEventTypes = auto()
InstanceStateEventTypes = auto()
NodePerformanceEventTypes = auto()
ControlPlaneEventTypes = auto()
DataPlaneEventTypes = auto()
TimerEventTypes = auto()
MLXEventTypes = auto()
PossibleEventOfEventTypeT = TypeVar("PossibleEventOfEventTypeT", bound=Enum)
# T=(A|B) <: U=(A|B|C) ==> Event[A|B] <: Event[A|BCategoryOfEventsT_cov = TypeVar(name="CategoryOfEventsT_cov", bound=EventCategories, covariant=True)
CategoryOfEventsT_cov = TypeVar(
name="CategoryOfEventsT_cov", bound=EventCategories, contravariant=True
)
CategoryOfEventsT_con = TypeVar(
name="CategoryOfEventsT_con", bound=EventCategories, contravariant=True
)
CategoryOfEventsT_inv = TypeVar(
name="CategoryOfEventsT_inv",
bound=EventCategories,
covariant=False,
contravariant=False,
# Here's the set of all possible events.
EventTypes = (
TaskEventTypes
| StreamingEventTypes
| InstanceEventTypes
| RunnerStatusEventTypes
| NodePerformanceEventTypes
| ControlPlaneEventTypes
| DataPlaneEventTypes
| TimerEventTypes
| TaskSagaEventTypes
)
class Event(BaseModel, Generic[PossibleEventOfEventTypeT]):
event_type: PossibleEventOfEventTypeT
event_category: EventCategories
check_event_type_union_is_consistent_with_registry(EVENT_TYPE_ENUMS, EventTypes)
class EventCategoryEnum(StrEnum):
MutatesTaskState = "MutatesTaskState"
MutatesTaskSagaState = "MutatesTaskSagaState"
MutatesRunnerStatus = "MutatesRunnerStatus"
MutatesInstanceState = "MutatesInstanceState"
MutatesNodePerformanceState = "MutatesNodePerformanceState"
MutatesControlPlaneState = "MutatesControlPlaneState"
MutatesDataPlaneState = "MutatesDataPlaneState"
EventCategory = (
Literal[EventCategoryEnum.MutatesControlPlaneState]
| Literal[EventCategoryEnum.MutatesTaskState]
| Literal[EventCategoryEnum.MutatesTaskSagaState]
| Literal[EventCategoryEnum.MutatesRunnerStatus]
| Literal[EventCategoryEnum.MutatesInstanceState]
| Literal[EventCategoryEnum.MutatesNodePerformanceState]
| Literal[EventCategoryEnum.MutatesDataPlaneState]
)
EventCategories = FrozenSet[EventCategory]
assert_literal_union_covers_enum(EventCategory, EventCategoryEnum)
class Event[SetMembersT: EventCategories | EventCategory](BaseModel):
event_type: EventTypes
event_category: SetMembersT
event_id: EventId
def check_origin_id(self, origin_id: NodeId) -> bool:
return True
def check_event_was_sent_by_correct_node(self, origin_id: NodeId) -> bool: ...
class TaskEvent(Event[TaskEventTypes]):
event_type: TaskEventTypes
class InstanceEvent(Event[InstanceEventTypes]):
event_type: InstanceEventTypes
class InstanceStateEvent(Event[InstanceStateEventTypes]):
event_type: InstanceStateEventTypes
class MLXEvent(Event[MLXEventTypes]):
event_type: MLXEventTypes
class NodePerformanceEvent(Event[NodePerformanceEventTypes]):
event_type: NodePerformanceEventTypes
class ControlPlaneEvent(Event[ControlPlaneEventTypes]):
event_type: ControlPlaneEventTypes
class StreamingEvent(Event[StreamingEventTypes]):
event_type: StreamingEventTypes
class DataPlaneEvent(Event[DataPlaneEventTypes]):
event_type: DataPlaneEventTypes
class TimerEvent(Event[TimerEventTypes]):
event_type: TimerEventTypes
class ResourceEvent(Event[ResourceEventTypes]):
event_type: ResourceEventTypes
class WrappedMessage(BaseModel, Generic[PossibleEventOfEventTypeT]):
message: Event[PossibleEventOfEventTypeT]
origin_id: NodeId
class EventFromEventLog[SetMembersT: EventCategories | EventCategory](BaseModel):
event: Event[SetMembersT]
origin: NodeId
idx_in_log: int = Field(gt=0)
@model_validator(mode="after")
def check_origin_id(self) -> "WrappedMessage[PossibleEventOfEventTypeT]":
if self.message.check_origin_id(self.origin_id):
def check_event_was_sent_by_correct_node(
self,
) -> "EventFromEventLog[SetMembersT]":
if self.event.check_event_was_sent_by_correct_node(self.origin):
return self
raise ValueError("Invalid Event: Origin ID Does Not Match")
class PersistedEvent(BaseModel, Generic[PossibleEventOfEventTypeT]):
event: Event[PossibleEventOfEventTypeT]
sequence_number: int = Field(gt=0)
def narrow_event_type[T: EventCategory, Q: EventCategories | EventCategory](
event: Event[Q],
target_category: T,
) -> Event[T]:
if target_category not in event.event_category:
raise ValueError(f"Event Does Not Contain Target Category {target_category}")
narrowed_event = event.model_copy(update={"event_category": {target_category}})
return cast(Event[T], narrowed_event)
class State(BaseModel, Generic[CategoryOfEventsT_cov]):
event_category: CategoryOfEventsT_cov
sequence_number: int = Field(default=0, ge=0)
def narrow_event_from_event_log_type[
T: EventCategory,
Q: EventCategories | EventCategory,
](
event: EventFromEventLog[Q],
target_category: T,
) -> EventFromEventLog[T]:
if target_category not in event.event.event_category:
raise ValueError(f"Event Does Not Contain Target Category {target_category}")
narrowed_event = event.model_copy(
update={"event": narrow_event_type(event.event, target_category)}
)
return cast(EventFromEventLog[T], narrowed_event)
AnnotatedEventType = Annotated[
Event[EventCategories], Field(discriminator="event_category")
class State[EventCategoryT: EventCategory](BaseModel):
event_category: EventCategoryT
last_event_applied_idx: int = Field(default=0, ge=0)
# Definitions for Type Variables
type Saga[EventCategoryT: EventCategory] = Callable[
[State[EventCategoryT], EventFromEventLog[EventCategoryT]],
Sequence[Event[EventCategories]],
]
EventTypeParser: TypeAdapter[AnnotatedEventType] = TypeAdapter(AnnotatedEventType)
# it's not possible to enforce this at compile time, so we have to do it at runtime
def mock_todo[T](something: T | None) -> T: ...
def apply(
state: State[CategoryOfEventsT_inv], event: Event[CategoryOfEventsT_inv]
) -> State[CategoryOfEventsT_inv]: ...
# T=(A|B) <: U=(A|B|C) ==> Apply[A|B] <: Apply[A|B|C]
SagaApplicator = Callable[
[State[CategoryOfEventsT_inv], Event[CategoryOfEventsT_inv]],
Sequence[Event[CategoryOfEventsT_inv]],
type Apply[EventCategoryT: EventCategory] = Callable[
[State[EventCategoryT], EventFromEventLog[EventCategoryT]],
State[EventCategoryT],
]
Saga = Callable[
[State[CategoryOfEventsT_inv], Event[CategoryOfEventsT_inv]],
Sequence[Event[CategoryOfEventsT_inv]],
class StateAndEvent[EventCategoryT: EventCategory](NamedTuple):
state: State[EventCategoryT]
event: EventFromEventLog[EventCategoryT]
type EffectHandler[EventCategoryT: EventCategory] = Callable[
[StateAndEvent[EventCategoryT], State[EventCategoryT]], None
]
Apply = Callable[
[State[CategoryOfEventsT_inv], Event[CategoryOfEventsT_inv]],
State[CategoryOfEventsT_inv],
type EventPublisher[EventCategoryT: EventCategory] = Callable[
[Event[EventCategoryT]], None
]
StateAndEvent = Tuple[State[CategoryOfEventsT_inv], Event[CategoryOfEventsT_inv]]
EffectHandler = Callable[
[StateAndEvent[CategoryOfEventsT_inv], State[CategoryOfEventsT_inv]], None
]
EventPublisher = Callable[[Event[CategoryOfEventsT_inv]], None]
class MutableState[EventCategoryT: EventCategories](Protocol):
def apply(
self,
event: Event[EventCategoryT],
applicator: Apply[EventCategoryT],
effect_handlers: Sequence[EffectHandler[EventCategoryT]],
) -> None: ...
class EventOutbox(Protocol):
# A component that can publish events
class EventPublisherProtocol(Protocol):
def send(self, events: Sequence[Event[EventCategories]]) -> None: ...
#
# T=[A|B] <: U=[A|B|C] => EventProcessor[A|B] :> EventProcessor[A|B|C]
#
class EventProcessor[EventCategoryT: EventCategories](Protocol):
# A component that can fetch events to apply
class EventFetcherProtocol[EventCategoryT: EventCategory](Protocol):
def get_events_to_apply(
self, state: State[EventCategoryT]
) -> Sequence[Event[EventCategoryT]]: ...
def get_saga_effect_handler[EventCategoryT: EventCategories](
# A component that can get the effect handler for a saga
def get_saga_effect_handler[EventCategoryT: EventCategory](
saga: Saga[EventCategoryT], event_publisher: EventPublisher[EventCategoryT]
) -> EffectHandler[EventCategoryT]:
def effect_handler(state_and_event: StateAndEvent[EventCategoryT]) -> None:
@@ -241,14 +234,16 @@ def get_saga_effect_handler[EventCategoryT: EventCategories](
return lambda state_and_event, _: effect_handler(state_and_event)
def get_effects_from_sagas[EventCategoryT: EventCategories](
def get_effects_from_sagas[EventCategoryT: EventCategory](
sagas: Sequence[Saga[EventCategoryT]],
event_publisher: EventPublisher[EventCategoryT],
) -> Sequence[EffectHandler[EventCategoryT]]:
return [get_saga_effect_handler(saga, event_publisher) for saga in sagas]
IdemKeyGenerator = Callable[[State[CategoryOfEventsT_cov], int], Sequence[EventId]]
type IdemKeyGenerator[EventCategoryT: EventCategory] = Callable[
[State[EventCategoryT], int], Sequence[EventId]
]
class CommandId(NewUUID):
@@ -261,14 +256,15 @@ class CommandTypes(str, Enum):
Delete = "Delete"
class Command[EventCategoryT: EventCategories, CommandType: CommandTypes](BaseModel):
class Command[
EventCategoryT: EventCategories | EventCategory,
CommandType: CommandTypes,
](BaseModel):
command_type: CommandType
command_id: CommandId
CommandTypeT = TypeVar("CommandTypeT", bound=CommandTypes, covariant=True)
Decide = Callable[
[State[CategoryOfEventsT_cov], Command[CategoryOfEventsT_cov, CommandTypeT]],
Sequence[Event[CategoryOfEventsT_cov]],
type Decide[EventCategoryT: EventCategory, CommandTypeT: CommandTypes] = Callable[
[State[EventCategoryT], Command[EventCategoryT, CommandTypeT]],
Sequence[Event[EventCategoryT]],
]

178
shared/types/events/events.py
View File

@@ -1,32 +1,21 @@
from __future__ import annotations
from typing import Any, Literal, Tuple
from pydantic import BaseModel
from typing import Literal, Tuple
from shared.types.common import NodeId
from shared.types.events.chunks import GenerationChunk
from shared.types.events.common import (
ControlPlaneEvent,
ControlPlaneEventTypes,
DataPlaneEvent,
DataPlaneEventTypes,
InstanceEvent,
Event,
EventCategoryEnum,
EventTypes,
InstanceEventTypes,
InstanceStateEvent,
InstanceStateEventTypes,
MLXEvent,
MLXEventTypes,
NodePerformanceEvent,
NodePerformanceEventTypes,
ResourceEvent,
ResourceEventTypes,
StreamingEvent,
RunnerStatusEventTypes,
StreamingEventTypes,
TaskEvent,
TaskEventTypes,
TimerEvent,
TimerEventTypes,
TimerId,
TaskSagaEventTypes,
)
from shared.types.networking.control_plane import (
ControlPlaneEdgeId,
@@ -37,149 +26,140 @@ from shared.types.networking.data_plane import (
DataPlaneEdgeId,
DataPlaneEdgeProfile,
)
from shared.types.profiling.common import NodePerformanceProfile, ProfiledResourceName
from shared.types.profiling.common import NodePerformanceProfile
from shared.types.tasks.common import (
TaskData,
TaskId,
TaskParams,
TaskState,
TaskStatusIncompleteType,
TaskStatusOtherType,
TaskStatusType,
TaskType,
)
from shared.types.worker.common import InstanceId, NodeStatus
from shared.types.worker.instances import InstanceData, InstanceStatus
from shared.types.worker.runners import RunnerId, RunnerState, RunnerStateType
from shared.types.worker.instances import InstanceParams, TypeOfInstance
from shared.types.worker.runners import RunnerId, RunnerStatus, RunnerStatusType
MLXEvent = Event[
frozenset(
{
EventCategoryEnum.MutatesTaskState,
EventCategoryEnum.MutatesControlPlaneState,
}
)
]
TaskEvent = Event[EventCategoryEnum.MutatesTaskState]
InstanceEvent = Event[EventCategoryEnum.MutatesInstanceState]
ControlPlaneEvent = Event[EventCategoryEnum.MutatesControlPlaneState]
DataPlaneEvent = Event[EventCategoryEnum.MutatesDataPlaneState]
NodePerformanceEvent = Event[EventCategoryEnum.MutatesNodePerformanceState]
class TimerData(BaseModel):
timer_id: TimerId
class TaskCreated[TaskTypeT: TaskType](TaskEvent):
event_type: TaskEventTypes = TaskEventTypes.TaskCreated
class TaskCreated(Event[EventCategoryEnum.MutatesTaskState]):
event_type: EventTypes = TaskEventTypes.TaskCreated
task_id: TaskId
task_data: TaskData
task_state: TaskState[Literal[TaskStatusIncompleteType.Pending], TaskTypeT]
task_params: TaskParams[TaskType]
task_state: TaskState[Literal[TaskStatusOtherType.Pending], TaskType]
on_instance: InstanceId
class TaskUpdated[TaskTypeT: TaskType](TaskEvent):
event_type: TaskEventTypes = TaskEventTypes.TaskUpdated
task_id: TaskId
update_data: TaskState[TaskStatusType, TaskTypeT]
class TaskDeleted(TaskEvent):
event_type: TaskEventTypes = TaskEventTypes.TaskDeleted
# Covers Cancellation Of Task, Non-Cancelled Tasks Perist
class TaskDeleted(Event[EventCategoryEnum.MutatesTaskState]):
event_type: EventTypes = TaskEventTypes.TaskDeleted
task_id: TaskId
class InstanceCreated(InstanceEvent):
event_type: InstanceEventTypes = InstanceEventTypes.InstanceCreated
class TaskStateUpdated(Event[EventCategoryEnum.MutatesTaskState]):
event_type: EventTypes = TaskEventTypes.TaskStateUpdated
task_state: TaskState[TaskStatusType, TaskType]
class InstanceCreated(Event[EventCategoryEnum.MutatesInstanceState]):
event_type: EventTypes = InstanceEventTypes.InstanceCreated
instance_id: InstanceId
instance_data: InstanceData
target_status: InstanceStatus
instance_params: InstanceParams
instance_type: TypeOfInstance
class InstanceDeleted(InstanceEvent):
event_type: InstanceEventTypes = InstanceEventTypes.InstanceDeleted
class InstanceActivated(Event[EventCategoryEnum.MutatesInstanceState]):
event_type: EventTypes = InstanceEventTypes.InstanceActivated
instance_id: InstanceId
class InstanceStatusUpdated(InstanceEvent):
event_type: InstanceEventTypes = InstanceEventTypes.InstanceStatusUpdated
class InstanceDeactivated(Event[EventCategoryEnum.MutatesInstanceState]):
event_type: EventTypes = InstanceEventTypes.InstanceDeactivated
instance_id: InstanceId
instance_status: InstanceStatus
class InstanceRunnerStateUpdated(InstanceStateEvent):
event_type: InstanceStateEventTypes = (
InstanceStateEventTypes.InstanceRunnerStateUpdated
)
class InstanceDeleted(Event[EventCategoryEnum.MutatesInstanceState]):
event_type: EventTypes = InstanceEventTypes.InstanceDeleted
instance_id: InstanceId
state_update: Tuple[RunnerId, RunnerState[RunnerStateType]]
class InstanceToBeReplacedAtomically(InstanceEvent):
event_type: InstanceEventTypes = InstanceEventTypes.InstanceToBeReplacedAtomically
transition: Tuple[InstanceId, InstanceId]
class InstanceReplacedAtomically(InstanceEvent):
event_type: InstanceEventTypes = InstanceEventTypes.InstanceReplacedAtomically
transition: Tuple[InstanceId, InstanceId]
class InstanceReplacedAtomically(Event[EventCategoryEnum.MutatesInstanceState]):
event_type: EventTypes = InstanceEventTypes.InstanceReplacedAtomically
instance_to_replace: InstanceId
new_instance_id: InstanceId
class MLXInferenceSagaPrepare(MLXEvent):
event_type: MLXEventTypes = MLXEventTypes.MLXInferenceSagaPrepare
class RunnerStatusUpdated(Event[EventCategoryEnum.MutatesRunnerStatus]):
event_type: EventTypes = RunnerStatusEventTypes.RunnerStatusUpdated
instance_id: InstanceId
state_update: Tuple[RunnerId, RunnerStatus[RunnerStatusType]]
class MLXInferenceSagaPrepare(Event[EventCategoryEnum.MutatesTaskSagaState]):
event_type: EventTypes = TaskSagaEventTypes.MLXInferenceSagaPrepare
task_id: TaskId
instance_id: InstanceId
class MLXInferenceSagaStartPrepare(MLXEvent):
event_type: MLXEventTypes = MLXEventTypes.MLXInferenceSagaStartPrepare
class MLXInferenceSagaStartPrepare(Event[EventCategoryEnum.MutatesTaskSagaState]):
event_type: EventTypes = TaskSagaEventTypes.MLXInferenceSagaStartPrepare
task_id: TaskId
instance_id: InstanceId
class NodePerformanceProfiled(NodePerformanceEvent):
event_type: NodePerformanceEventTypes = (
NodePerformanceEventTypes.NodePerformanceProfiled
)
class NodePerformanceMeasured(Event[EventCategoryEnum.MutatesNodePerformanceState]):
event_type: EventTypes = NodePerformanceEventTypes.NodePerformanceMeasured
node_id: NodeId
node_profile: NodePerformanceProfile
class WorkerConnected(ControlPlaneEvent):
event_type: ControlPlaneEventTypes = ControlPlaneEventTypes.WorkerConnected
class WorkerConnected(Event[EventCategoryEnum.MutatesControlPlaneState]):
event_type: EventTypes = ControlPlaneEventTypes.WorkerConnected
edge: DataPlaneEdge
class WorkerStatusUpdated(ControlPlaneEvent):
event_type: ControlPlaneEventTypes = ControlPlaneEventTypes.WorkerStatusUpdated
class WorkerStatusUpdated(Event[EventCategoryEnum.MutatesControlPlaneState]):
event_type: EventTypes = ControlPlaneEventTypes.WorkerStatusUpdated
node_id: NodeId
node_state: NodeStatus
class WorkerDisconnected(ControlPlaneEvent):
event_type: ControlPlaneEventTypes = ControlPlaneEventTypes.WorkerConnected
class WorkerDisconnected(Event[EventCategoryEnum.MutatesControlPlaneState]):
event_type: EventTypes = ControlPlaneEventTypes.WorkerConnected
vertex_id: ControlPlaneEdgeId
class ChunkGenerated(StreamingEvent):
event_type: StreamingEventTypes = StreamingEventTypes.ChunkGenerated
class ChunkGenerated(Event[EventCategoryEnum.MutatesTaskState]):
event_type: EventTypes = StreamingEventTypes.ChunkGenerated
task_id: TaskId
instance_id: InstanceId
chunk: Any
chunk: GenerationChunk
class DataPlaneEdgeCreated(DataPlaneEvent):
event_type: DataPlaneEventTypes = DataPlaneEventTypes.DataPlaneEdgeCreated
class DataPlaneEdgeCreated(Event[EventCategoryEnum.MutatesDataPlaneState]):
event_type: EventTypes = DataPlaneEventTypes.DataPlaneEdgeCreated
vertex: ControlPlaneEdgeType
class DataPlaneEdgeProfiled(DataPlaneEvent):
event_type: DataPlaneEventTypes = DataPlaneEventTypes.DataPlaneEdgeProfiled
class DataPlaneEdgeReplacedAtomically(Event[EventCategoryEnum.MutatesDataPlaneState]):
event_type: EventTypes = DataPlaneEventTypes.DataPlaneEdgeReplacedAtomically
edge_id: DataPlaneEdgeId
edge_profile: DataPlaneEdgeProfile
class DataPlaneEdgeDeleted(DataPlaneEvent):
event_type: DataPlaneEventTypes = DataPlaneEventTypes.DataPlaneEdgeDeleted
class DataPlaneEdgeDeleted(Event[EventCategoryEnum.MutatesDataPlaneState]):
event_type: EventTypes = DataPlaneEventTypes.DataPlaneEdgeDeleted
edge_id: DataPlaneEdgeId
class TimerScheduled(TimerEvent):
event_type: TimerEventTypes = TimerEventTypes.TimerCreated
timer_data: TimerData
class TimerFired(TimerEvent):
event_type: TimerEventTypes = TimerEventTypes.TimerFired
timer_data: TimerData
class ResourceProfiled(ResourceEvent):
event_type: ResourceEventTypes = ResourceEventTypes.ResourceProfiled
resource_name: ProfiledResourceName
resource_profile: NodePerformanceProfile

131
shared/types/events/registry.py Normal file
View File

@@ -0,0 +1,131 @@
from types import UnionType
from typing import Annotated, Any, Mapping, Type, get_args
from pydantic import Field, TypeAdapter
from shared.constants import EXO_ERROR_REPORTING_MESSAGE
from shared.types.events.common import (
ControlPlaneEventTypes,
DataPlaneEventTypes,
Event,
EventCategories,
EventTypes,
InstanceEventTypes,
NodePerformanceEventTypes,
RunnerStatusEventTypes,
StreamingEventTypes,
TaskEventTypes,
TaskSagaEventTypes,
)
from shared.types.events.events import (
ChunkGenerated,
DataPlaneEdgeCreated,
DataPlaneEdgeDeleted,
DataPlaneEdgeReplacedAtomically,
InstanceCreated,
InstanceDeleted,
InstanceReplacedAtomically,
MLXInferenceSagaPrepare,
MLXInferenceSagaStartPrepare,
NodePerformanceMeasured,
RunnerStatusUpdated,
TaskCreated,
TaskDeleted,
TaskStateUpdated,
WorkerConnected,
WorkerDisconnected,
WorkerStatusUpdated,
)
"""
class EventTypeNames(StrEnum):
TaskEventType = auto()
InstanceEvent = auto()
NodePerformanceEvent = auto()
ControlPlaneEvent = auto()
StreamingEvent = auto()
DataPlaneEvent = auto()
TimerEvent = auto()
MLXEvent = auto()
check_event_categories_are_defined_for_all_event_types(EVENT_TYPE_ENUMS, EventTypeNames)
"""
EventRegistry: Mapping[EventTypes, Type[Any]] = {
TaskEventTypes.TaskCreated: TaskCreated,
TaskEventTypes.TaskStateUpdated: TaskStateUpdated,
TaskEventTypes.TaskDeleted: TaskDeleted,
InstanceEventTypes.InstanceCreated: InstanceCreated,
InstanceEventTypes.InstanceDeleted: InstanceDeleted,
InstanceEventTypes.InstanceReplacedAtomically: InstanceReplacedAtomically,
RunnerStatusEventTypes.RunnerStatusUpdated: RunnerStatusUpdated,
NodePerformanceEventTypes.NodePerformanceMeasured: NodePerformanceMeasured,
ControlPlaneEventTypes.WorkerConnected: WorkerConnected,
ControlPlaneEventTypes.WorkerStatusUpdated: WorkerStatusUpdated,
ControlPlaneEventTypes.WorkerDisconnected: WorkerDisconnected,
StreamingEventTypes.ChunkGenerated: ChunkGenerated,
DataPlaneEventTypes.DataPlaneEdgeCreated: DataPlaneEdgeCreated,
DataPlaneEventTypes.DataPlaneEdgeReplacedAtomically: DataPlaneEdgeReplacedAtomically,
DataPlaneEventTypes.DataPlaneEdgeDeleted: DataPlaneEdgeDeleted,
TaskSagaEventTypes.MLXInferenceSagaPrepare: MLXInferenceSagaPrepare,
TaskSagaEventTypes.MLXInferenceSagaStartPrepare: MLXInferenceSagaStartPrepare,
}
# Sanity Check.
def check_registry_has_all_event_types() -> None:
event_types: tuple[EventTypes, ...] = get_args(EventTypes)
missing_event_types = set(event_types) - set(EventRegistry.keys())
assert not missing_event_types, (
f"{EXO_ERROR_REPORTING_MESSAGE()}"
f"There's an event missing from the registry: {missing_event_types}"
)
def check_union_of_all_events_is_consistent_with_registry(
registry: Mapping[EventTypes, Type[Any]], union_type: UnionType
) -> None:
type_of_each_registry_entry = set(type(event_type) for event_type in registry)
type_of_each_entry_in_union = set(get_args(union_type))
missing_from_union = type_of_each_registry_entry - type_of_each_entry_in_union
assert not missing_from_union, (
f"{EXO_ERROR_REPORTING_MESSAGE()}"
f"Event classes in registry are missing from all_events union: {missing_from_union}"
)
extra_in_union = type_of_each_entry_in_union - type_of_each_registry_entry
assert not extra_in_union, (
f"{EXO_ERROR_REPORTING_MESSAGE()}"
f"Event classes in all_events union are missing from registry: {extra_in_union}"
)
AllEvents = (
TaskCreated
| TaskStateUpdated
| TaskDeleted
| InstanceCreated
| InstanceDeleted
| InstanceReplacedAtomically
| RunnerStatusUpdated
| NodePerformanceMeasured
| WorkerConnected
| WorkerStatusUpdated
| WorkerDisconnected
| ChunkGenerated
| DataPlaneEdgeCreated
| DataPlaneEdgeReplacedAtomically
| DataPlaneEdgeDeleted
| MLXInferenceSagaPrepare
| MLXInferenceSagaStartPrepare
)
# Run the sanity check
check_union_of_all_events_is_consistent_with_registry(EventRegistry, AllEvents)
_EventType = Annotated[AllEvents, Field(discriminator="event_type")]
EventParser: TypeAdapter[Event[EventCategories]] = TypeAdapter(_EventType)

68
shared/types/events/sanity_checking.py Normal file
View File

@@ -0,0 +1,68 @@
from enum import Enum, StrEnum
from types import UnionType
from typing import Any, LiteralString, Sequence, Set, Type, get_args
from shared.constants import EXO_ERROR_REPORTING_MESSAGE
def check_event_type_union_is_consistent_with_registry(
event_type_enums: Sequence[Type[Enum]], event_types: UnionType
) -> None:
"""Assert that every enum value from _EVENT_TYPE_ENUMS satisfies EventTypes."""
event_types_inferred_from_union = set(get_args(event_types))
event_types_inferred_from_registry = [
member for enum_class in event_type_enums for member in enum_class
]
# Check that each registry value belongs to one of the types in the union
for tag_of_event_type in event_types_inferred_from_registry:
event_type = type(tag_of_event_type)
assert event_type in event_types_inferred_from_union, (
f"{EXO_ERROR_REPORTING_MESSAGE()}"
f"There's a mismatch between the registry of event types and the union of possible event types."
f"The enum value {tag_of_event_type} for type {event_type} is not covered by {event_types_inferred_from_union}."
)
def check_event_categories_are_defined_for_all_event_types(
event_definitions: Sequence[Type[Enum]], event_categories: Type[StrEnum]
) -> None:
"""Assert that the event category names are consistent with the event type enums."""
expected_category_tags: list[str] = [
enum_class.__name__ for enum_class in event_definitions
]
tag_of_event_categories: list[str] = list(event_categories.__members__.values())
assert tag_of_event_categories == expected_category_tags, (
f"{EXO_ERROR_REPORTING_MESSAGE()}"
f"The values of the enum EventCategories are not named after the event type enums."
f"These are the missing categories: {set(expected_category_tags) - set(tag_of_event_categories)}"
f"These are the extra categories: {set(tag_of_event_categories) - set(expected_category_tags)}"
)
def assert_literal_union_covers_enum[TEnum: StrEnum](
literal_union: UnionType,
enum_type: Type[TEnum],
) -> None:
enum_values: Set[Any] = {member.value for member in enum_type}
def _flatten(tp: UnionType) -> Set[Any]:
values: Set[Any] = set()
args: tuple[LiteralString, ...] = get_args(tp)
for arg in args:
payloads: tuple[TEnum, ...] = get_args(arg)
for payload in payloads:
values.add(payload.value)
return values
literal_values: Set[Any] = _flatten(literal_union)
assert enum_values == literal_values, (
f"{EXO_ERROR_REPORTING_MESSAGE()}"
f"The values of the enum {enum_type} are not covered by the literal union {literal_union}.\n"
f"These are the missing values: {enum_values - literal_values}\n"
f"These are the extra values: {literal_values - enum_values}\n"
)

11
shared/types/graphs/common.py
View File

@@ -41,8 +41,8 @@ class Edge(
class GraphData(BaseModel, Generic[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]):
edges: Mapping[EdgeIdT, EdgeData[EdgeTypeT]]
vertices: Mapping[VertexIdT, VertexData[VertexTypeT]]
edges: Mapping[EdgeIdT, EdgeData[EdgeTypeT]] = {}
vertices: Mapping[VertexIdT, VertexData[VertexTypeT]] = {}
class GraphProtocol(Protocol, Generic[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]):
@@ -111,11 +111,12 @@ class MutableGraphProtocol(GraphProtocol[EdgeTypeT, VertexTypeT, EdgeIdT, Vertex
class Graph(
BaseModel,
Generic[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
GraphProtocol[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
MutableGraphProtocol[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
):
graph_data: GraphData[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]
graph_data: GraphData[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT] = GraphData[
EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT
]()
# the first element in the return value is the filtered graph; the second is the

3
shared/types/models/metadata.py
View File

@@ -1,8 +1,9 @@
from typing import Annotated
from typing import Annotated, final
from pydantic import BaseModel, PositiveInt
@final
class ModelMetadata(BaseModel):
pretty_name: str
storage_size_kilobytes: Annotated[int, PositiveInt]

14
shared/types/models/sources.py
View File

@@ -1,27 +1,26 @@
from enum import Enum
from typing import Annotated, Any, Generic, Literal, TypeVar, Union, final
from typing import Annotated, Any, Literal, Union, final
from pydantic import AnyHttpUrl, BaseModel, Field, TypeAdapter
from shared.types.models.common import ModelId
@final
class SourceType(str, Enum):
HuggingFace = "HuggingFace"
GitHub = "GitHub"
@final
class SourceFormatType(str, Enum):
HuggingFaceTransformers = "HuggingFaceTransformers"
T = TypeVar("T", bound=SourceType)
S = TypeVar("S", bound=SourceFormatType)
RepoPath = Annotated[str, Field(pattern=r"^[^/]+/[^/]+$")]
class BaseModelSource(BaseModel, Generic[T, S]):
class BaseModelSource[T: SourceType, S: SourceFormatType](BaseModel):
model_uuid: ModelId
source_type: T
source_format: S
@@ -50,15 +49,16 @@ class HuggingFaceModelSource(
@final
class GitHubModelSource(BaseModelSource[SourceType.GitHub, S]):
class GitHubModelSource(BaseModelSource[SourceType.GitHub, SourceFormatType]):
source_type: Literal[SourceType.GitHub] = SourceType.GitHub
source_format: SourceFormatType
source_data: GitHubModelSourceData
_ModelSource = Annotated[
Union[
HuggingFaceModelSource,
GitHubModelSource[SourceFormatType.HuggingFaceTransformers],
GitHubModelSource,
],
Field(discriminator="source_type"),
]

45
shared/types/networking/topology.py
View File

@@ -1,5 +1,5 @@
from shared.graphs.networkx import NetworkXGraph
from shared.types.common import NodeId
from shared.types.graphs.common import Graph, GraphData
from shared.types.networking.control_plane import ControlPlaneEdgeId
from shared.types.networking.data_plane import (
DataPlaneEdgeData,
@@ -9,64 +9,37 @@ from shared.types.worker.common import NodeStatus
class DataPlaneTopology(
Graph[
NetworkXGraph[
DataPlaneEdgeData,
None,
DataPlaneEdgeId,
NodeId,
]
):
graph_data: GraphData[
DataPlaneEdgeData,
None,
DataPlaneEdgeId,
NodeId,
]
pass
class OrphanedPartOfDataPlaneTopology(
Graph[
NetworkXGraph[
DataPlaneEdgeData,
None,
DataPlaneEdgeId,
NodeId,
]
):
graph_data: GraphData[
DataPlaneEdgeData,
None,
DataPlaneEdgeId,
NodeId,
]
pass
class ControlPlaneTopology(
Graph[
None,
NodeStatus,
ControlPlaneEdgeId,
NodeId,
]
):
graph_data: GraphData[
None,
NodeStatus,
ControlPlaneEdgeId,
NodeId,
]
class ControlPlaneTopology(NetworkXGraph[None, NodeStatus, ControlPlaneEdgeId, NodeId]):
pass
class OrphanedPartOfControlPlaneTopology(
Graph[
NetworkXGraph[
None,
NodeStatus,
ControlPlaneEdgeId,
NodeId,
]
):
graph_data: GraphData[
None,
NodeStatus,
ControlPlaneEdgeId,
NodeId,
]
pass

59
shared/types/states/master.py
View File

@@ -1,19 +1,21 @@
from collections.abc import Mapping, Sequence
from enum import Enum
from queue import Queue
from typing import Generic, TypeVar
from typing import Generic, Literal, TypeVar
from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter
from shared.types.common import NodeId
from shared.types.events.common import (
Event,
EventCategories,
EventCategory,
EventCategoryEnum,
State,
)
from shared.types.graphs.resource_graph import ResourceGraph
from shared.types.networking.data_plane import (
DataPlaneEdge,
DataPlaneEdgeAdapter,
DataPlaneEdgeId,
)
from shared.types.networking.topology import (
@@ -24,8 +26,9 @@ from shared.types.networking.topology import (
)
from shared.types.profiling.common import NodePerformanceProfile
from shared.types.states.shared import SharedState
from shared.types.tasks.common import TaskData
from shared.types.worker.instances import InstanceData, InstanceId
from shared.types.tasks.common import TaskParams, TaskType
from shared.types.worker.common import NodeStatus
from shared.types.worker.instances import InstanceId, InstanceParams
class ExternalCommand(BaseModel): ...
@@ -42,44 +45,56 @@ class CachePolicy(BaseModel, Generic[CachePolicyTypeT]):
policy_type: CachePolicyTypeT
class NodePerformanceProfileState(State[EventCategories.NodePerformanceEventTypes]):
class NodePerformanceProfileState(State[EventCategoryEnum.MutatesNodePerformanceState]):
node_profiles: Mapping[NodeId, NodePerformanceProfile]
class DataPlaneNetworkState(State[EventCategories.DataPlaneEventTypes]):
topology: DataPlaneTopology
history: Sequence[OrphanedPartOfDataPlaneTopology]
class DataPlaneNetworkState(State[EventCategoryEnum.MutatesDataPlaneState]):
event_category: Literal[EventCategoryEnum.MutatesDataPlaneState] = (
EventCategoryEnum.MutatesDataPlaneState
)
topology: DataPlaneTopology = DataPlaneTopology(
edge_base=DataPlaneEdgeAdapter, vertex_base=TypeAdapter(None)
)
history: Sequence[OrphanedPartOfDataPlaneTopology] = []
def delete_edge(self, edge_id: DataPlaneEdgeId) -> None: ...
def add_edge(self, edge: DataPlaneEdge) -> None: ...
class ControlPlaneNetworkState(State[EventCategories.ControlPlaneEventTypes]):
topology: ControlPlaneTopology
history: Sequence[OrphanedPartOfControlPlaneTopology]
class ControlPlaneNetworkState(State[EventCategoryEnum.MutatesControlPlaneState]):
event_category: Literal[EventCategoryEnum.MutatesControlPlaneState] = (
EventCategoryEnum.MutatesControlPlaneState
)
topology: ControlPlaneTopology = ControlPlaneTopology(
edge_base=TypeAdapter(None), vertex_base=TypeAdapter(NodeStatus)
)
history: Sequence[OrphanedPartOfControlPlaneTopology] = []
def delete_edge(self, edge_id: DataPlaneEdgeId) -> None: ...
def add_edge(self, edge: DataPlaneEdge) -> None: ...
class MasterState(SharedState):
data_plane_network_state: DataPlaneNetworkState
control_plane_network_state: ControlPlaneNetworkState
job_inbox: Queue[TaskData]
job_outbox: Queue[TaskData]
cache_policy: CachePolicy[CachePolicyType]
data_plane_network_state: DataPlaneNetworkState = DataPlaneNetworkState()
control_plane_network_state: ControlPlaneNetworkState = ControlPlaneNetworkState()
job_inbox: Queue[TaskParams[TaskType]] = Queue()
job_outbox: Queue[TaskParams[TaskType]] = Queue()
cache_policy: CachePolicy[CachePolicyType] = CachePolicy[CachePolicyType](
policy_type=CachePolicyType.KeepAll
)
def get_shard_assignments(
inbox: Queue[ExternalCommand],
outbox: Queue[ExternalCommand],
resource_graph: ResourceGraph,
current_instances: Mapping[InstanceId, InstanceData],
current_instances: Mapping[InstanceId, InstanceParams],
cache_policy: CachePolicy[CachePolicyType],
) -> Mapping[InstanceId, InstanceData]: ...
) -> Mapping[InstanceId, InstanceParams]: ...
def get_transition_events(
current_instances: Mapping[InstanceId, InstanceData],
target_instances: Mapping[InstanceId, InstanceData],
) -> Sequence[Event[EventCategories]]: ...
current_instances: Mapping[InstanceId, InstanceParams],
target_instances: Mapping[InstanceId, InstanceParams],
) -> Sequence[Event[EventCategory]]: ...

48
shared/types/states/shared.py
View File

@@ -1,27 +1,55 @@
from collections.abc import Mapping
from typing import Sequence
from typing import Literal, Sequence
from pydantic import BaseModel
from shared.types.common import NodeId
from shared.types.events.common import EventCategories, State
from shared.types.tasks.common import Task, TaskId, TaskStatusType, TaskType
from shared.types.events.common import EventCategoryEnum, State
from shared.types.tasks.common import (
Task,
TaskId,
TaskSagaEntry,
TaskStatusType,
TaskType,
)
from shared.types.worker.common import InstanceId
from shared.types.worker.instances import BaseInstance
from shared.types.worker.runners import RunnerId, RunnerStatus, RunnerStatusType
class KnownInstances(State[EventCategories.InstanceStateEventTypes]):
instances: Mapping[InstanceId, BaseInstance]
class Instances(State[EventCategoryEnum.MutatesInstanceState]):
event_category: Literal[EventCategoryEnum.MutatesInstanceState] = (
EventCategoryEnum.MutatesInstanceState
)
instances: Mapping[InstanceId, BaseInstance] = {}
class Tasks(State[EventCategories.TaskEventTypes]):
tasks: Mapping[TaskId, Task[TaskType, TaskStatusType]]
class Tasks(State[EventCategoryEnum.MutatesTaskState]):
event_category: Literal[EventCategoryEnum.MutatesTaskState] = (
EventCategoryEnum.MutatesTaskState
)
tasks: Mapping[TaskId, Task[TaskType, TaskStatusType]] = {}
class TaskSagas(State[EventCategoryEnum.MutatesTaskSagaState]):
event_category: Literal[EventCategoryEnum.MutatesTaskSagaState] = (
EventCategoryEnum.MutatesTaskSagaState
)
task_sagas: Mapping[TaskId, Sequence[TaskSagaEntry]] = {}
class Runners(State[EventCategoryEnum.MutatesRunnerStatus]):
event_category: Literal[EventCategoryEnum.MutatesRunnerStatus] = (
EventCategoryEnum.MutatesRunnerStatus
)
runner_statuses: Mapping[RunnerId, RunnerStatus[RunnerStatusType]] = {}
class SharedState(BaseModel):
node_id: NodeId
known_instances: KnownInstances
compute_tasks: Tasks
instances: Instances = Instances()
runners: Runners = Runners()
tasks: Tasks = Tasks()
task_sagas: TaskSagas = TaskSagas()
def get_node_id(self) -> NodeId: ...

4
shared/types/states/worker.py
View File

@@ -2,14 +2,14 @@ from collections.abc import Mapping
from shared.types.common import NodeId
from shared.types.events.common import (
EventCategories,
EventCategory,
State,
)
from shared.types.states.shared import SharedState
from shared.types.worker.common import NodeStatus
class NodeStatusState(State[EventCategories.ControlPlaneEventTypes]):
class NodeStatusState(State[EventCategory.MutatesControlPlaneState]):
node_status: Mapping[NodeId, NodeStatus]

125
shared/types/tasks/common.py
View File

@@ -1,17 +1,18 @@
from collections.abc import Mapping
from enum import Enum
from typing import Annotated, Any, Generic, Literal, TypeVar, Union
from typing import Annotated, Generic, Literal, TypeVar, Union, final
import openai.types.chat as openai
from pydantic import BaseModel, Field, TypeAdapter
from shared.types.common import NewUUID
from shared.types.worker.common import InstanceId, RunnerId
from shared.types.worker.common import InstanceId
class TaskId(NewUUID):
pass
@final
class TaskType(str, Enum):
ChatCompletionNonStreaming = "ChatCompletionNonStreaming"
ChatCompletionStreaming = "ChatCompletionStreaming"
@@ -20,123 +21,69 @@ class TaskType(str, Enum):
TaskTypeT = TypeVar("TaskTypeT", bound=TaskType, covariant=True)
class BaseTaskData(BaseModel, Generic[TaskTypeT]):
task_type: TaskTypeT
class TaskParams(BaseModel, Generic[TaskTypeT]): ...
# Custom message types that mirror OpenAI's but are designed for serialization
class ChatCompletionMessage(BaseModel):
role: Literal["system", "user", "assistant", "developer", "tool", "function"]
content: str | None = None
name: str | None = None
tool_calls: list[dict[str, Any]] | None = None
tool_call_id: str | None = None
function_call: dict[str, Any] | None = None
class ChatCompletionParams(BaseModel):
model: str
messages: list[ChatCompletionMessage]
frequency_penalty: float | None = None
logit_bias: dict[str, int] | None = None
logprobs: bool | None = None
top_logprobs: int | None = None
max_tokens: int | None = None
n: int | None = None
presence_penalty: float | None = None
response_format: dict[str, Any] | None = None
seed: int | None = None
stop: str | list[str] | None = None
stream: bool = False
temperature: float | None = None
top_p: float | None = None
tools: list[dict[str, Any]] | None = None
tool_choice: str | dict[str, Any] | None = None
parallel_tool_calls: bool | None = None
user: str | None = None
class ChatCompletionNonStreamingTask(BaseTaskData[TaskType.ChatCompletionNonStreaming]):
@final
class ChatCompletionNonStreamingTask(TaskParams[TaskType.ChatCompletionNonStreaming]):
task_type: Literal[TaskType.ChatCompletionNonStreaming] = (
TaskType.ChatCompletionNonStreaming
)
task_data: ChatCompletionParams
task_data: openai.completion_create_params.CompletionCreateParams
class ChatCompletionStreamingTask(BaseTaskData[TaskType.ChatCompletionStreaming]):
@final
class ChatCompletionStreamingTask(TaskParams[TaskType.ChatCompletionStreaming]):
task_type: Literal[TaskType.ChatCompletionStreaming] = (
TaskType.ChatCompletionStreaming
)
task_data: ChatCompletionParams
task_data: openai.completion_create_params.CompletionCreateParams
TaskData = Annotated[
ChatCompletionNonStreamingTask | ChatCompletionStreamingTask,
Field(discriminator="task_type"),
]
TaskDataValidator: TypeAdapter[TaskData] = TypeAdapter(TaskData)
class TaskStatusIncompleteType(str, Enum):
Pending = "Pending"
Running = "Running"
@final
class TaskStatusFailedType(str, Enum):
Failed = "Failed"
@final
class TaskStatusCompleteType(str, Enum):
Complete = "Complete"
TaskStatusType = Union[TaskStatusIncompleteType, TaskStatusCompleteType]
@final
class TaskStatusOtherType(str, Enum):
Pending = "Pending"
Running = "Running"
TaskStatusType = TaskStatusCompleteType | TaskStatusFailedType | TaskStatusOtherType
class TaskArtifact[TaskTypeT: TaskType, TaskStatusTypeT: TaskStatusType](BaseModel): ...
class IncompleteTaskArtifact[TaskTypeT: TaskType](
TaskArtifact[TaskTypeT, TaskStatusIncompleteType]
):
@final
class NoTaskArtifact[TaskTypeT: TaskType](TaskArtifact[TaskTypeT, TaskStatusOtherType]):
pass
class TaskStatusUpdate[TaskStatusTypeT: TaskStatusType](BaseModel):
task_status: TaskStatusTypeT
class PendingTaskStatus(TaskStatusUpdate[TaskStatusIncompleteType.Pending]):
task_status: Literal[TaskStatusIncompleteType.Pending] = (
TaskStatusIncompleteType.Pending
)
class RunningTaskStatus(TaskStatusUpdate[TaskStatusIncompleteType.Running]):
task_status: Literal[TaskStatusIncompleteType.Running] = (
TaskStatusIncompleteType.Running
)
class CompletedTaskStatus(TaskStatusUpdate[TaskStatusCompleteType.Complete]):
task_status: Literal[TaskStatusCompleteType.Complete] = (
TaskStatusCompleteType.Complete
)
class FailedTaskStatus(TaskStatusUpdate[TaskStatusIncompleteType.Failed]):
task_status: Literal[TaskStatusIncompleteType.Failed] = (
TaskStatusIncompleteType.Failed
)
error_message: Mapping[RunnerId, str]
@final
class FailedTaskArtifact[TaskTypeT: TaskType](
TaskArtifact[TaskTypeT, TaskStatusFailedType]
):
error_message: str
@final
class TaskState[TaskStatusTypeT: TaskStatusType, TaskTypeT: TaskType](BaseModel):
task_status: TaskStatusUpdate[TaskStatusTypeT]
task_status: TaskStatusTypeT
task_artifact: TaskArtifact[TaskTypeT, TaskStatusTypeT]
class BaseTask[TaskTypeT: TaskType, TaskStatusTypeT: TaskStatusType](BaseModel):
task_type: TaskTypeT
task_data: TaskData
task_state: TaskState[TaskStatusTypeT, TaskTypeT]
task_params: TaskParams[TaskTypeT]
task_stats: TaskState[TaskStatusTypeT, TaskTypeT]
on_instance: InstanceId
@@ -148,11 +95,17 @@ BaseTaskAnnotated = Annotated[
Field(discriminator="task_type"),
]
BaseTaskValidator: TypeAdapter[BaseTask[TaskType, TaskStatusType]] = TypeAdapter(
BaseTaskParser: TypeAdapter[BaseTask[TaskType, TaskStatusType]] = TypeAdapter(
BaseTaskAnnotated
)
class TaskSagaEntry(BaseModel):
task_id: TaskId
instance_id: InstanceId
@final
class Task[TaskTypeT: TaskType, TaskStatusTypeT: TaskStatusType](
BaseTask[TaskTypeT, TaskStatusTypeT]
):

10
shared/types/worker/commands_runner.py
View File

@@ -4,11 +4,9 @@ from typing import Annotated, Generic, Literal, TypeVar
from pydantic import BaseModel, Field, TypeAdapter
from shared.openai import FinishReason
# Accept TaskData so the runner can handle both streaming and non-streaming chat tasks.
from shared.types.tasks.common import TaskData
from shared.types.api import ChatTask
from shared.types.worker.mlx import Host
from shared.types.worker.shards import ShardMeta
from shared.types.worker.shards import PartitionStrategy, ShardMetadata
## Messages passed TO the runner
@@ -28,7 +26,7 @@ class BaseRunnerMessage(BaseModel, Generic[MT]):
class SetupMessage(BaseRunnerMessage[MessageType.Setup]):
type: Literal[MessageType.Setup] = Field(default=MessageType.Setup, frozen=True)
model_shard_meta: ShardMeta
model_shard_meta: ShardMetadata[PartitionStrategy]
hosts: list[Host]
@@ -36,7 +34,7 @@ class ChatTaskMessage(BaseRunnerMessage[MessageType.ChatTask]):
type: Literal[MessageType.ChatTask] = Field(
default=MessageType.ChatTask, frozen=True
)
task: TaskData
task: ChatTask
class ExitMessage(BaseRunnerMessage[MessageType.Exit]):

11
shared/types/worker/downloads.py
View File

@@ -2,11 +2,9 @@ from enum import Enum
from typing import (
Annotated,
Callable,
Generic,
Literal,
NewType,
Sequence,
TypeVar,
Union,
)
@@ -15,7 +13,7 @@ from pydantic import BaseModel, Field, PositiveInt
from shared.types.common import NodeId
from shared.types.models.common import ModelId
from shared.types.models.sources import ModelSource
from shared.types.worker.shards import BaseShardMeta, PartitionStrategy
from shared.types.worker.shards import PartitionStrategy, ShardMetadata
class DownloadProgressData(BaseModel):
@@ -30,10 +28,7 @@ class DownloadStatus(str, Enum):
Failed = "Failed"
DownloadStatusT = TypeVar("DownloadStatusT", bound=DownloadStatus)
class BaseDownloadProgress(BaseModel, Generic[DownloadStatusT]):
class BaseDownloadProgress[DownloadStatusT: DownloadStatus](BaseModel):
node_id: NodeId
download_status: DownloadStatusT
@@ -80,6 +75,6 @@ DownloadEffectHandler = Callable[
def download_shard(
model_id: ModelId,
model_source: ModelSource,
shard_meta: BaseShardMeta[PartitionStrategy],
shard_metadata: ShardMetadata[PartitionStrategy],
effect_handlers: Sequence[DownloadEffectHandler],
) -> None: ...

17
shared/types/worker/instances.py
View File

@@ -1,34 +1,25 @@
from collections.abc import Mapping
from enum import Enum
from pydantic import BaseModel
from shared.types.worker.common import InstanceId
from shared.types.worker.runners import (
RunnerId,
RunnerState,
RunnerStateType,
ShardAssignments,
)
class InstanceStatus(str, Enum):
class TypeOfInstance(str, Enum):
ACTIVE = "active"
INACTIVE = "inactive"
class InstanceState(BaseModel):
runner_states: Mapping[RunnerId, RunnerState[RunnerStateType]]
class InstanceData(BaseModel):
class InstanceParams(BaseModel):
shard_assignments: ShardAssignments
class BaseInstance(BaseModel):
instance_data: InstanceData
instance_state: InstanceState
instance_status: InstanceStatus
instance_params: InstanceParams
instance_type: TypeOfInstance
class Instance(BaseInstance):

6
shared/types/worker/mlx.py
View File

@@ -7,11 +7,7 @@ class Host(BaseModel):
port: int
@field_validator("port")
@classmethod
def check_port(cls, v: int) -> int:
def check_port(self, v: int) -> int:
if not (0 <= v <= 65535):
raise ValueError("Port must be between 0 and 65535")
return v
def __str__(self) -> str:
return f"{self.host}:{self.port}"

39
shared/types/worker/resource_monitor.py
View File

@@ -1,9 +1,8 @@
import asyncio
from abc import ABC, abstractmethod
from collections.abc import Coroutine
from typing import Callable, Set
from typing import Callable, List, Set
from shared.types.events.events import ResourceProfiled
from shared.types.profiling.common import (
MemoryPerformanceProfile,
NodePerformanceProfile,
@@ -11,58 +10,44 @@ from shared.types.profiling.common import (
)
class EventLog:
def append(self, event: ResourceProfiled) -> None: ...
class ResourceCollector(ABC):
"""
Details a single resource (or resource type) that is being monitored by the resource monitor.
"""
def __init__(self, name: str):
self.name = name
name = str
@abstractmethod
async def collect(self) -> NodePerformanceProfile: ...
class SystemResourceCollector(ResourceCollector):
def __init__(self):
super().__init__("system")
name = "system"
@abstractmethod
async def collect(self) -> SystemPerformanceProfile: ...
class MemoryResourceCollector(ResourceCollector):
def __init__(self):
super().__init__("memory")
name = "memory"
@abstractmethod
async def collect(self) -> MemoryPerformanceProfile: ...
class ResourceMonitor:
def __init__(
self,
collectors: list[ResourceCollector],
effect_handlers: Set[Callable[[NodePerformanceProfile], None]],
):
self.effect_handlers: Set[Callable[[NodePerformanceProfile], None]] = (
effect_handlers
)
self.collectors: list[ResourceCollector] = collectors
data_collectors: List[ResourceCollector]
effect_handlers: Set[Callable[[NodePerformanceProfile], None]]
# Since there's no implementation, this breaks the typechecker.
# self.collectors: list[ResourceCollector] = [
# SystemResourceCollector(),
# MemoryResourceCollector(),
# ]
# Since there's no implementation, this breaks the typechecker.
# self.collectors: list[ResourceCollector] = [
# SystemResourceCollector(),
# MemoryResourceCollector(),
# ]
async def _collect(self) -> list[NodePerformanceProfile]:
tasks: list[Coroutine[None, None, NodePerformanceProfile]] = [
collector.collect() for collector in self.collectors
collector.collect() for collector in self.data_collectors
]
return await asyncio.gather(*tasks)

67
shared/types/worker/runners.py
View File

@@ -1,29 +1,17 @@
from collections.abc import Mapping, Sequence
from enum import Enum
from typing import Generic, Literal, TypeVar
from typing import Annotated, Generic, Literal, TypeVar
from pydantic import BaseModel, model_validator
from pydantic import BaseModel, Field, TypeAdapter, model_validator
from shared.types.common import NodeId
from shared.types.models.common import ModelId
from shared.types.worker.common import RunnerId
from shared.types.worker.downloads import BaseDownloadProgress, DownloadStatus
from shared.types.worker.shards import BaseShardMeta, PartitionStrategy
from shared.types.worker.shards import PartitionStrategy, ShardMetadata
class RunnerError(Exception):
error_type: str
error_message: str
traceback: str
def __init__(self, error_type: str, error_message: str, traceback: str):
self.error_type = error_type
self.error_message = error_message
self.traceback = traceback
super().__init__(f"{error_type}: {error_message}\n{traceback}")
class RunnerStateType(str, Enum):
class RunnerStatusType(str, Enum):
Rejected = "Rejected"
Starting = "Starting"
Downloading = "Downloading"
@@ -31,48 +19,51 @@ class RunnerStateType(str, Enum):
Failed = "Failed"
RunnerStateTypeT = TypeVar("RunnerStateTypeT", bound=RunnerStateType)
RunnerStatusTypeT = TypeVar("RunnerStatusTypeT", bound=RunnerStatusType)
class RunnerState(BaseModel, Generic[RunnerStateTypeT]):
runner_state: RunnerStateTypeT
class RunnerStatus(BaseModel, Generic[RunnerStatusTypeT]):
runner_status: RunnerStatusTypeT
class RejectedRunnerState(RunnerState[RunnerStateType.Rejected]):
runner_state: Literal[RunnerStateType.Rejected]
class RejectedRunnerStatus(RunnerStatus[RunnerStatusType.Rejected]):
runner_status: Literal[RunnerStatusType.Rejected]
class StartingRunnerState(RunnerState[RunnerStateType.Starting]):
runner_state: Literal[RunnerStateType.Starting]
class StartingRunnerStatus(RunnerStatus[RunnerStatusType.Starting]):
runner_status: Literal[RunnerStatusType.Starting]
class DownloadingRunnerState(RunnerState[RunnerStateType.Downloading]):
runner_state: Literal[RunnerStateType.Downloading]
class DownloadingRunnerStatus(RunnerStatus[RunnerStatusType.Downloading]):
runner_status: Literal[RunnerStatusType.Downloading]
download_progress: BaseDownloadProgress[DownloadStatus]
class RunningRunnerState(RunnerState[RunnerStateType.Running]):
runner_state: Literal[RunnerStateType.Running]
class RunningRunnerStatus(RunnerStatus[RunnerStatusType.Running]):
runner_status: Literal[RunnerStatusType.Running]
class FailedRunnerState(RunnerState[RunnerStateType.Failed]):
runner_state: Literal[RunnerStateType.Failed]
class FailedRunnerStatus(RunnerStatus[RunnerStatusType.Failed]):
runner_status: Literal[RunnerStatusType.Failed]
error_message: str | None = None
class RunnerData(BaseModel):
runner_id: RunnerId
runner_state: RunnerState[RunnerStateType] = RunnerState(
runner_state=RunnerStateType.Starting
)
PartitionStrategyT = TypeVar(name="PartitionStrategyT", bound=PartitionStrategy)
_RunnerStatus = Annotated[
RejectedRunnerStatus
| StartingRunnerStatus
| DownloadingRunnerStatus
| RunningRunnerStatus
| FailedRunnerStatus,
Field,
]
RunnerStatusParser: TypeAdapter[RunnerStatus[RunnerStatusType]] = TypeAdapter(
_RunnerStatus
)
class ShardAssignments(BaseModel):
model_id: ModelId
runner_to_shard: Mapping[RunnerId, BaseShardMeta[PartitionStrategy]]
runner_to_shard: Mapping[RunnerId, ShardMetadata[PartitionStrategy]]
node_to_runner: Mapping[NodeId, Sequence[RunnerId]]
@model_validator(mode="after")

30
shared/types/worker/shards.py
View File

@@ -1,7 +1,7 @@
from enum import Enum
from typing import Annotated, Generic, Literal, TypeVar
from typing import Annotated, Literal
from pydantic import BaseModel, ConfigDict, DirectoryPath, Field, TypeAdapter
from pydantic import BaseModel, DirectoryPath, Field, TypeAdapter
from shared.types.common import NodeId
from shared.types.models.common import ModelId
@@ -11,44 +11,44 @@ class PartitionStrategy(str, Enum):
pipeline = "pipeline"
PartitionStrategyT = TypeVar(name="PartitionStrategyT", bound=PartitionStrategy)
class BaseShardMeta(BaseModel, Generic[PartitionStrategyT]):
class ShardMetadata[PartitionStrategyT: PartitionStrategy](BaseModel):
"""
Defines a specific shard of the model that is ready to be run on a device.
Replaces previous `Shard` object.
"""
partition_strategy: PartitionStrategyT
device_rank: int
world_size: int
model_id: ModelId
model_path: DirectoryPath
class PipelineShardMeta(BaseShardMeta[Literal[PartitionStrategy.pipeline]]):
class PipelineShardMetadata(ShardMetadata[PartitionStrategy.pipeline]):
"""
Pipeline parallelism shard meta.
"""
model_config = ConfigDict(use_enum_values=False)
partition_strategy: Literal[PartitionStrategy.pipeline] = PartitionStrategy.pipeline
partition_strategy: Literal[PartitionStrategy.pipeline] = Field(
default=PartitionStrategy.pipeline, frozen=True
)
start_layer: Annotated[int, Field(ge=0)]
end_layer: Annotated[int, Field(ge=0)]
_ShardMeta = Annotated[PipelineShardMeta, Field(discriminator="partition_strategy")]
ShardMeta = _ShardMeta # Public alias for the discriminated union
ShardMetaAdapter: TypeAdapter[BaseShardMeta[PartitionStrategy]] = TypeAdapter(
_ShardMeta
_ShardMetadata = Annotated[
PipelineShardMetadata, Field(discriminator="partition_strategy")
]
ShardMetaParser: TypeAdapter[ShardMetadata[PartitionStrategy]] = TypeAdapter(
_ShardMetadata
)
class ShardPlacement(BaseModel, Generic[PartitionStrategyT]):
class ShardPlacement[PartitionStrategyT: PartitionStrategy](BaseModel):
"""
A shard placement is the description of a model distributed across a set of nodes.
The Generic[PartitionStrategyT] enforces that the shard assignments all use the same partition strategy.
"""
model_id: ModelId
shard_assignments: dict[NodeId, BaseShardMeta[PartitionStrategyT]]
shard_assignments: dict[NodeId, ShardMetadata[PartitionStrategyT]]

190
uv.lock generated
View File

@@ -42,45 +42,13 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/a1/ee/48ca1a7c89ffec8b6a0c5d02b89c305671d5ffd8d3c94acf8b8c408575bb/anyio-4.9.0-py3-none-any.whl", hash = "sha256:9f76d541cad6e36af7beb62e978876f3b41e3e04f2c1fbf0884604c0a9c4d93c", size = 100916, upload-time = "2025-03-17T00:02:52.713Z" },
]
[[package]]
name = "basedpyright"
version = "1.30.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "nodejs-wheel-binaries", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/1d/d8/a2c9dfa97de316fe228c978bc4677cadb4dc44971d52db026405b8e58377/basedpyright-1.30.0.tar.gz", hash = "sha256:45f5c94b92a8cb9506998c6d29129becd5a2118f14fdbc0df289b96d6a8ff8bc", size = 22059435, upload-time = "2025-07-09T12:12:58.642Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/f6/62/65a06c403ac5e7fc0e11b5ab7617a584786a9606c4a19b7269dcc3c61eb3/basedpyright-1.30.0-py3-none-any.whl", hash = "sha256:782afca88f88a24429a82d900a77deafe88ac88af256774ee304528dd93344f2", size = 11537772, upload-time = "2025-07-09T12:12:54.568Z" },
]
[[package]]
name = "certifi"
version = "2025.7.9"
version = "2025.6.15"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/de/8a/c729b6b60c66a38f590c4e774decc4b2ec7b0576be8f1aa984a53ffa812a/certifi-2025.7.9.tar.gz", hash = "sha256:c1d2ec05395148ee10cf672ffc28cd37ea0ab0d99f9cc74c43e588cbd111b079", size = 160386, upload-time = "2025-07-09T02:13:58.874Z" }
sdist = { url = "https://files.pythonhosted.org/packages/73/f7/f14b46d4bcd21092d7d3ccef689615220d8a08fb25e564b65d20738e672e/certifi-2025.6.15.tar.gz", hash = "sha256:d747aa5a8b9bbbb1bb8c22bb13e22bd1f18e9796defa16bab421f7f7a317323b", size = 158753, upload-time = "2025-06-15T02:45:51.329Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/66/f3/80a3f974c8b535d394ff960a11ac20368e06b736da395b551a49ce950cce/certifi-2025.7.9-py3-none-any.whl", hash = "sha256:d842783a14f8fdd646895ac26f719a061408834473cfc10203f6a575beb15d39", size = 159230, upload-time = "2025-07-09T02:13:57.007Z" },
]
[[package]]
name = "charset-normalizer"
version = "3.4.2"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/e4/33/89c2ced2b67d1c2a61c19c6751aa8902d46ce3dacb23600a283619f5a12d/charset_normalizer-3.4.2.tar.gz", hash = "sha256:5baececa9ecba31eff645232d59845c07aa030f0c81ee70184a90d35099a0e63", size = 126367, upload-time = "2025-05-02T08:34:42.01Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/ea/12/a93df3366ed32db1d907d7593a94f1fe6293903e3e92967bebd6950ed12c/charset_normalizer-3.4.2-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:926ca93accd5d36ccdabd803392ddc3e03e6d4cd1cf17deff3b989ab8e9dbcf0", size = 199622, upload-time = "2025-05-02T08:32:56.363Z" },
{ url = "https://files.pythonhosted.org/packages/04/93/bf204e6f344c39d9937d3c13c8cd5bbfc266472e51fc8c07cb7f64fcd2de/charset_normalizer-3.4.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:eba9904b0f38a143592d9fc0e19e2df0fa2e41c3c3745554761c5f6447eedabf", size = 143435, upload-time = "2025-05-02T08:32:58.551Z" },
{ url = "https://files.pythonhosted.org/packages/22/2a/ea8a2095b0bafa6c5b5a55ffdc2f924455233ee7b91c69b7edfcc9e02284/charset_normalizer-3.4.2-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:3fddb7e2c84ac87ac3a947cb4e66d143ca5863ef48e4a5ecb83bd48619e4634e", size = 153653, upload-time = "2025-05-02T08:33:00.342Z" },
{ url = "https://files.pythonhosted.org/packages/b6/57/1b090ff183d13cef485dfbe272e2fe57622a76694061353c59da52c9a659/charset_normalizer-3.4.2-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:98f862da73774290f251b9df8d11161b6cf25b599a66baf087c1ffe340e9bfd1", size = 146231, upload-time = "2025-05-02T08:33:02.081Z" },
{ url = "https://files.pythonhosted.org/packages/e2/28/ffc026b26f441fc67bd21ab7f03b313ab3fe46714a14b516f931abe1a2d8/charset_normalizer-3.4.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c9379d65defcab82d07b2a9dfbfc2e95bc8fe0ebb1b176a3190230a3ef0e07c", size = 148243, upload-time = "2025-05-02T08:33:04.063Z" },
{ url = "https://files.pythonhosted.org/packages/c0/0f/9abe9bd191629c33e69e47c6ef45ef99773320e9ad8e9cb08b8ab4a8d4cb/charset_normalizer-3.4.2-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e635b87f01ebc977342e2697d05b56632f5f879a4f15955dfe8cef2448b51691", size = 150442, upload-time = "2025-05-02T08:33:06.418Z" },
{ url = "https://files.pythonhosted.org/packages/67/7c/a123bbcedca91d5916c056407f89a7f5e8fdfce12ba825d7d6b9954a1a3c/charset_normalizer-3.4.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:1c95a1e2902a8b722868587c0e1184ad5c55631de5afc0eb96bc4b0d738092c0", size = 145147, upload-time = "2025-05-02T08:33:08.183Z" },
{ url = "https://files.pythonhosted.org/packages/ec/fe/1ac556fa4899d967b83e9893788e86b6af4d83e4726511eaaad035e36595/charset_normalizer-3.4.2-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:ef8de666d6179b009dce7bcb2ad4c4a779f113f12caf8dc77f0162c29d20490b", size = 153057, upload-time = "2025-05-02T08:33:09.986Z" },
{ url = "https://files.pythonhosted.org/packages/2b/ff/acfc0b0a70b19e3e54febdd5301a98b72fa07635e56f24f60502e954c461/charset_normalizer-3.4.2-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:32fc0341d72e0f73f80acb0a2c94216bd704f4f0bce10aedea38f30502b271ff", size = 156454, upload-time = "2025-05-02T08:33:11.814Z" },
{ url = "https://files.pythonhosted.org/packages/92/08/95b458ce9c740d0645feb0e96cea1f5ec946ea9c580a94adfe0b617f3573/charset_normalizer-3.4.2-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:289200a18fa698949d2b39c671c2cc7a24d44096784e76614899a7ccf2574b7b", size = 154174, upload-time = "2025-05-02T08:33:13.707Z" },
{ url = "https://files.pythonhosted.org/packages/78/be/8392efc43487ac051eee6c36d5fbd63032d78f7728cb37aebcc98191f1ff/charset_normalizer-3.4.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:4a476b06fbcf359ad25d34a057b7219281286ae2477cc5ff5e3f70a246971148", size = 149166, upload-time = "2025-05-02T08:33:15.458Z" },
{ url = "https://files.pythonhosted.org/packages/20/94/c5790835a017658cbfabd07f3bfb549140c3ac458cfc196323996b10095a/charset_normalizer-3.4.2-py3-none-any.whl", hash = "sha256:7f56930ab0abd1c45cd15be65cc741c28b1c9a34876ce8c17a2fa107810c0af0", size = 52626, upload-time = "2025-05-02T08:34:40.053Z" },
{ url = "https://files.pythonhosted.org/packages/84/ae/320161bd181fc06471eed047ecce67b693fd7515b16d495d8932db763426/certifi-2025.6.15-py3-none-any.whl", hash = "sha256:2e0c7ce7cb5d8f8634ca55d2ba7e6ec2689a2fd6537d8dec1296a477a4910057", size = 157650, upload-time = "2025-06-15T02:45:49.977Z" },
]
[[package]]
@@ -108,7 +76,6 @@ darwin = [
[package.dev-dependencies]
dev = [
{ name = "basedpyright", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "maturin", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "pytest", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "pytest-asyncio", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -125,7 +92,6 @@ provides-extras = ["darwin"]
[package.metadata.requires-dev]
dev = [
{ name = "basedpyright", specifier = ">=1.29.4" },
{ name = "maturin", specifier = ">=1.9.0" },
{ name = "pytest", specifier = ">=8.4.0" },
{ name = "pytest-asyncio", specifier = ">=1.0.0" },
@@ -143,10 +109,14 @@ version = "0.1.0"
source = { editable = "master" }
dependencies = [
{ name = "exo-shared", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "fastapi", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
]
[package.metadata]
requires-dist = [{ name = "exo-shared", editable = "shared" }]
requires-dist = [
{ name = "exo-shared", editable = "shared" },
{ name = "fastapi", specifier = ">=0.116.0" },
]
[[package]]
name = "exo-networking"
@@ -158,11 +128,13 @@ name = "exo-shared"
version = "0.1.0"
source = { editable = "shared" }
dependencies = [
{ name = "networkx", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "openai", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "pathlib", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "pydantic", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "rich", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "rustworkx", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
]
[package.dev-dependencies]
@@ -172,11 +144,13 @@ dev = [
[package.metadata]
requires-dist = [
{ name = "networkx", specifier = ">=3.5" },
{ name = "openai", specifier = ">=1.93.0" },
{ name = "pathlib", specifier = ">=1.0.1" },
{ name = "protobuf", specifier = ">=6.31.1" },
{ name = "pydantic", specifier = ">=2.11.7" },
{ name = "rich", specifier = ">=14.0.0" },
{ name = "rustworkx", specifier = ">=0.16.0" },
]
[package.metadata.requires-dev]
@@ -296,6 +270,66 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/76/c6/c88e154df9c4e1a2a66ccf0005a88dfb2650c1dffb6f5ce603dfbd452ce3/idna-3.10-py3-none-any.whl", hash = "sha256:946d195a0d259cbba61165e88e65941f16e9b36ea6ddb97f00452bae8b1287d3", size = 70442, upload-time = "2024-09-15T18:07:37.964Z" },
]
[[package]]
name = "fastapi"
version = "0.116.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "pydantic", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "starlette", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/20/38/e1da78736143fd885c36213a3ccc493c384ae8fea6a0f0bc272ef42ebea8/fastapi-0.116.0.tar.gz", hash = "sha256:80dc0794627af0390353a6d1171618276616310d37d24faba6648398e57d687a", size = 296518, upload-time = "2025-07-07T15:09:27.82Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/2f/68/d80347fe2360445b5f58cf290e588a4729746e7501080947e6cdae114b1f/fastapi-0.116.0-py3-none-any.whl", hash = "sha256:fdcc9ed272eaef038952923bef2b735c02372402d1203ee1210af4eea7a78d2b", size = 95625, upload-time = "2025-07-07T15:09:26.348Z" },
]
[[package]]
name = "h11"
version = "0.16.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/01/ee/02a2c011bdab74c6fb3c75474d40b3052059d95df7e73351460c8588d963/h11-0.16.0.tar.gz", hash = "sha256:4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff1", size = 101250, upload-time = "2025-04-24T03:35:25.427Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/04/4b/29cac41a4d98d144bf5f6d33995617b185d14b22401f75ca86f384e87ff1/h11-0.16.0-py3-none-any.whl", hash = "sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86", size = 37515, upload-time = "2025-04-24T03:35:24.344Z" },
]
[[package]]
name = "httpcore"
version = "1.0.9"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "certifi", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "h11", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/06/94/82699a10bca87a5556c9c59b5963f2d039dbd239f25bc2a63907a05a14cb/httpcore-1.0.9.tar.gz", hash = "sha256:6e34463af53fd2ab5d807f399a9b45ea31c3dfa2276f15a2c3f00afff6e176e8", size = 85484, upload-time = "2025-04-24T22:06:22.219Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55", size = 78784, upload-time = "2025-04-24T22:06:20.566Z" },
]
[[package]]
name = "httpx"
version = "0.28.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "anyio", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "certifi", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "httpcore", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "idna", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/b1/df/48c586a5fe32a0f01324ee087459e112ebb7224f646c0b5023f5e79e9956/httpx-0.28.1.tar.gz", hash = "sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc", size = 141406, upload-time = "2024-12-06T15:37:23.222Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" },
]
[[package]]
name = "idna"
version = "3.10"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/f1/70/7703c29685631f5a7590aa73f1f1d3fa9a380e654b86af429e0934a32f7d/idna-3.10.tar.gz", hash = "sha256:12f65c9b470abda6dc35cf8e63cc574b1c52b11df2c86030af0ac09b01b13ea9", size = 190490, upload-time = "2024-09-15T18:07:39.745Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/76/c6/c88e154df9c4e1a2a66ccf0005a88dfb2650c1dffb6f5ce603dfbd452ce3/idna-3.10-py3-none-any.whl", hash = "sha256:946d195a0d259cbba61165e88e65941f16e9b36ea6ddb97f00452bae8b1287d3", size = 70442, upload-time = "2024-09-15T18:07:37.964Z" },
]
[[package]]
name = "iniconfig"
version = "2.1.0"
@@ -305,18 +339,6 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/2c/e1/e6716421ea10d38022b952c159d5161ca1193197fb744506875fbb87ea7b/iniconfig-2.1.0-py3-none-any.whl", hash = "sha256:9deba5723312380e77435581c6bf4935c94cbfab9b1ed33ef8d238ea168eb760", size = 6050, upload-time = "2025-03-19T20:10:01.071Z" },
]
[[package]]
name = "jinja2"
version = "3.1.6"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "markupsafe", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" },
]
[[package]]
name = "jiter"
version = "0.10.0"
@@ -440,17 +462,12 @@ wheels = [
]
[[package]]
name = "nodejs-wheel-binaries"
version = "22.17.0"
name = "networkx"
version = "3.5"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/d3/86/8962d1d24ff480f4dd31871f42c8e0d8e2c851cd558a07ee689261d310ab/nodejs_wheel_binaries-22.17.0.tar.gz", hash = "sha256:529142012fb8fd20817ef70e2ef456274df4f49933292e312c8bbc7285af6408", size = 8068, upload-time = "2025-06-29T20:24:25.002Z" }
sdist = { url = "https://files.pythonhosted.org/packages/6c/4f/ccdb8ad3a38e583f214547fd2f7ff1fc160c43a75af88e6aec213404b96a/networkx-3.5.tar.gz", hash = "sha256:d4c6f9cf81f52d69230866796b82afbccdec3db7ae4fbd1b65ea750feed50037", size = 2471065, upload-time = "2025-05-29T11:35:07.804Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/5d/53/b942c6da4ff6f87a315033f6ff6fed8fd3c22047d7ff5802badaa5dfc2c2/nodejs_wheel_binaries-22.17.0-py2.py3-none-macosx_11_0_arm64.whl", hash = "sha256:6545a6f6d2f736d9c9e2eaad7e599b6b5b2d8fd4cbd2a1df0807cbcf51b9d39b", size = 51003554, upload-time = "2025-06-29T20:23:47.042Z" },
{ url = "https://files.pythonhosted.org/packages/e2/b7/7184a9ad2364912da22f2fe021dc4a3301721131ef7759aeb4a1f19db0b4/nodejs_wheel_binaries-22.17.0-py2.py3-none-macosx_11_0_x86_64.whl", hash = "sha256:4bea5b994dd87c20f8260031ea69a97c3d282e2d4472cc8908636a313a830d00", size = 51936848, upload-time = "2025-06-29T20:23:52.064Z" },
{ url = "https://files.pythonhosted.org/packages/e9/7a/0ea425147b8110b8fd65a6c21cfd3bd130cdec7766604361429ef870d799/nodejs_wheel_binaries-22.17.0-py2.py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:885508615274a22499dd5314759c1cf96ba72de03e6485d73b3e5475e7f12662", size = 57925230, upload-time = "2025-06-29T20:23:56.81Z" },
{ url = "https://files.pythonhosted.org/packages/23/5f/10a3f2ac08a839d065d9ccfd6d9df66bc46e100eaf87a8a5cf149eb3fb8e/nodejs_wheel_binaries-22.17.0-py2.py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:90f38ce034a602bcab534d55cbe0390521e73e5dcffdd1c4b34354b932172af2", size = 58457829, upload-time = "2025-06-29T20:24:01.945Z" },
{ url = "https://files.pythonhosted.org/packages/ed/a4/d2ca331e16eef0974eb53702df603c54f77b2a7e2007523ecdbf6cf61162/nodejs_wheel_binaries-22.17.0-py2.py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:5eed087855b644c87001fe04036213193963ccd65e7f89949e9dbe28e7743d9b", size = 59778054, upload-time = "2025-06-29T20:24:07.14Z" },
{ url = "https://files.pythonhosted.org/packages/be/2b/04e0e7f7305fe2ba30fd4610bfb432516e0f65379fe6c2902f4b7b1ad436/nodejs_wheel_binaries-22.17.0-py2.py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:715f413c81500f0770ea8936ef1fc2529b900da8054cbf6da67cec3ee308dc76", size = 60830079, upload-time = "2025-06-29T20:24:12.21Z" },
{ url = "https://files.pythonhosted.org/packages/eb/8d/776adee7bbf76365fdd7f2552710282c79a4ead5d2a46408c9043a2b70ba/networkx-3.5-py3-none-any.whl", hash = "sha256:0030d386a9a06dee3565298b4a734b68589749a544acbb6c412dc9e2489ec6ec", size = 2034406, upload-time = "2025-05-29T11:35:04.961Z" },
]
[[package]]
@@ -479,7 +496,7 @@ wheels = [
[[package]]
name = "openai"
version = "1.93.3"
version = "1.93.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "anyio", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -491,9 +508,9 @@ dependencies = [
{ name = "tqdm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/e0/66/fadc0cad6a229c6a85c3aa5f222a786ec4d9bf14c2a004f80ffa21dbaf21/openai-1.93.3.tar.gz", hash = "sha256:488b76399238c694af7e4e30c58170ea55e6f65038ab27dbe95b5077a00f8af8", size = 487595, upload-time = "2025-07-09T14:08:27.789Z" }
sdist = { url = "https://files.pythonhosted.org/packages/e4/d7/e91c6a9cf71726420cddf539852ee4c29176ebb716a702d9118d0409fd8e/openai-1.93.0.tar.gz", hash = "sha256:988f31ade95e1ff0585af11cc5a64510225e4f5cd392698c675d0a9265b8e337", size = 486573, upload-time = "2025-06-27T21:21:39.421Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/8b/b9/0df6351b25c6bd494c534d2a8191dc9460fb5bb09c88b1427775d49fde05/openai-1.93.3-py3-none-any.whl", hash = "sha256:41aaa7594c7d141b46eed0a58dcd75d20edcc809fdd2c931ecbb4957dc98a892", size = 755132, upload-time = "2025-07-09T14:08:25.533Z" },
{ url = "https://files.pythonhosted.org/packages/64/46/a10d9df4673df56f71201d129ba1cb19eaff3366d08c8664d61a7df52e65/openai-1.93.0-py3-none-any.whl", hash = "sha256:3d746fe5498f0dd72e0d9ab706f26c91c0f646bf7459e5629af8ba7c9dbdf090", size = 755038, upload-time = "2025-06-27T21:21:37.532Z" },
]
[[package]]
@@ -790,6 +807,55 @@ sentencepiece = [
{ name = "sentencepiece", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
]
[[package]]
name = "rustworkx"
version = "0.16.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/a5/c4/6d6ef39e57610d54c5f106dc3dece9eebce8b9d52d561ae092e3aede1b66/rustworkx-0.16.0.tar.gz", hash = "sha256:9f0dcb83f38d5ca2c3a683eb9b6951c8aec3262fbfe5141946a7ee5ba37e0bb6", size = 349524, upload-time = "2025-01-24T01:22:34.686Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/f8/70/36f5916aee41ffe4f604ad75742eb1bb1b849fb568e010555f9d159cd93e/rustworkx-0.16.0-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:476a6c67b0142acd941691943750cc6737a48372304489969c2b62d30aaf4c27", size = 2141999, upload-time = "2025-01-24T01:21:50.3Z" },
{ url = "https://files.pythonhosted.org/packages/94/47/7e7c37fb73efcc87be6414b235534605c4008a4cdbd92a61db23b878eecd/rustworkx-0.16.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:bef2ef42870f806af93979b457e240f6dfa4f867ca33965c620f3a804409ed3a", size = 1940309, upload-time = "2025-01-24T01:21:52.053Z" },
{ url = "https://files.pythonhosted.org/packages/c6/42/a6d6b3137be55ef1d887becdf6b64b0917c7d437bd483065a88500a55603/rustworkx-0.16.0-cp39-abi3-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0db3a73bf68b3e66c08322a2fc95d3aa663d037d9b4e49c3509da4898d3529cc", size = 2195350, upload-time = "2025-01-24T01:21:53.785Z" },
{ url = "https://files.pythonhosted.org/packages/59/d2/1bc99df831c132c4b7420a85ce9150e065f4c993798f31b6a4229f238398/rustworkx-0.16.0-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4f12a13d7486234fa2a84746d5e41f436bf9df43548043e7a232f48804ff8c61", size = 1971689, upload-time = "2025-01-24T17:09:26.338Z" },
{ url = "https://files.pythonhosted.org/packages/b5/3b/1125e7eb834f4408bcec3cee79947efd504c715fb7ab1876f8cd4bbca497/rustworkx-0.16.0-cp39-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:89efd5c3a4653ddacc55ca39f28b261d43deec7d678f8f8fc6b76b5087f1dfea", size = 3297342, upload-time = "2025-01-24T03:18:48.885Z" },
{ url = "https://files.pythonhosted.org/packages/4f/e2/e21187b255c6211d71db0d08a44fc16771038b2af41712d66c408d9bec16/rustworkx-0.16.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ec0c12aac8c54910ace20ac6ada4b890cd39f95f69100514715f8ad7af9041e4", size = 2110107, upload-time = "2025-01-24T01:21:58.884Z" },
{ url = "https://files.pythonhosted.org/packages/3c/79/e3fcff21f31253ea85ef196bf2fcabad7802b11468f7d3a5d592cd0ac789/rustworkx-0.16.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:d650e39fc1a1534335f7517358ebfc3478bb235428463cfcd7c5750d50377b33", size = 2007544, upload-time = "2025-01-26T04:16:53.807Z" },
{ url = "https://files.pythonhosted.org/packages/67/04/741ed09c2b0dc0f360f85270c1179ed433785372ac9ab6ab26d3dd3ae02d/rustworkx-0.16.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:293180b83509ee9bff4c3af7ccc1024f6528d61b65d0cb7320bd31924f10cb71", size = 2172787, upload-time = "2025-01-24T01:22:01.282Z" },
]
[[package]]
name = "sniffio"
version = "1.3.1"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/a2/87/a6771e1546d97e7e041b6ae58d80074f81b7d5121207425c964ddf5cfdbd/sniffio-1.3.1.tar.gz", hash = "sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc", size = 20372, upload-time = "2024-02-25T23:20:04.057Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2", size = 10235, upload-time = "2024-02-25T23:20:01.196Z" },
]
[[package]]
name = "starlette"
version = "0.46.2"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "anyio", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/ce/20/08dfcd9c983f6a6f4a1000d934b9e6d626cff8d2eeb77a89a68eef20a2b7/starlette-0.46.2.tar.gz", hash = "sha256:7f7361f34eed179294600af672f565727419830b54b7b084efe44bb82d2fccd5", size = 2580846, upload-time = "2025-04-13T13:56:17.942Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/8b/0c/9d30a4ebeb6db2b25a841afbb80f6ef9a854fc3b41be131d249a977b4959/starlette-0.46.2-py3-none-any.whl", hash = "sha256:595633ce89f8ffa71a015caed34a5b2dc1c0cdb3f0f1fbd1e69339cf2abeec35", size = 72037, upload-time = "2025-04-13T13:56:16.21Z" },
]
[[package]]
name = "tqdm"
version = "4.67.1"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/a8/4b/29b4ef32e036bb34e4ab51796dd745cdba7ed47ad142a9f4a1eb8e0c744d/tqdm-4.67.1.tar.gz", hash = "sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2", size = 169737, upload-time = "2024-11-24T20:12:22.481Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl", hash = "sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2", size = 78540, upload-time = "2024-11-24T20:12:19.698Z" },
]
[[package]]
name = "types-protobuf"
version = "6.30.2.20250703"

13
worker/logging.py Normal file
View File

@@ -0,0 +1,13 @@
from collections.abc import Set
from typing import Literal
from shared.logging.common import LogEntry, LogEntryType
class WorkerUninitialized(LogEntry[Literal["master_uninitialized"]]):
entry_destination: Set[LogEntryType] = {LogEntryType.cluster}
entry_type: Literal["master_uninitialized"] = "master_uninitialized"
message: str = "No master state found, creating new one."
WorkerLogEntries = WorkerUninitialized

6
worker/main.py
View File

@@ -1,6 +0,0 @@
def main():
print("Hello from worker!")
if __name__ == "__main__":
main()