← back to Exo
Merge Basic Interfaces
b0bd95100501ef2f46091370510f12ae99c290d7 · 2025-07-09 19:04:21 +0100 · Arbion Halili
Co-authored-by: Alex Cheema <alexcheema123@gmail.com>
Co-authored-by: Seth Howes <sethshowes@gmail.com>
Co-authored-by: Matt Beton <matthew.beton@gmail.com>
Co-authored-by: Andrei Cravtov <the.andrei.cravtov@gmail.com>
Files touched
R100 .github/actions/conditional-commit/conditional-commit.yml .github/actions/conditional-commit/action.ymlA .github/actions/lint-check/action.ymlA .github/actions/verify-clean/action.ymlM .github/workflows/pipeline.ymlM flake.nixM justfileA master/api.pyM master/idempotency.pyM shared/logger.pyA shared/openai.pyA shared/protobufs/types/mlx/nn/__init__.pyiM shared/pyproject.tomlA shared/types/api.pyA shared/types/common.pyD shared/types/event_sourcing.pyA shared/types/events/chunks.pyA shared/types/events/common.pyA shared/types/events/events.pyA shared/types/graphs/common.pyA shared/types/graphs/resource_graph.pyA shared/types/models/common.pyA shared/types/models/metadata.pyA shared/types/models/model.pyA shared/types/models/sources.pyA shared/types/networking/control_plane.pyA shared/types/networking/data_plane.pyA shared/types/networking/services.pyA shared/types/networking/topology.pyA shared/types/profiling/common.pyA shared/types/states/master.pyA shared/types/states/shared.pyA shared/types/states/worker.pyA shared/types/tasks/common.pyA shared/types/worker/commands_runner.pyA shared/types/worker/common.pyA shared/types/worker/downloads.pyA shared/types/worker/instances.pyA shared/types/worker/mlx.pyA shared/types/worker/resource_monitor.pyA shared/types/worker/runners.pyA shared/types/worker/shards.pyA shared/utils.pyM uv.lock
Diff
commit b0bd95100501ef2f46091370510f12ae99c290d7
Author: Arbion Halili <99731180+ToxicPine@users.noreply.github.com>
Date: Wed Jul 9 19:04:21 2025 +0100
Merge Basic Interfaces
Co-authored-by: Alex Cheema <alexcheema123@gmail.com>
Co-authored-by: Seth Howes <sethshowes@gmail.com>
Co-authored-by: Matt Beton <matthew.beton@gmail.com>
Co-authored-by: Andrei Cravtov <the.andrei.cravtov@gmail.com>
---
.../{conditional-commit.yml => action.yml} | 0
.github/actions/lint-check/action.yml | 10 +
.github/actions/verify-clean/action.yml | 20 ++
.github/workflows/pipeline.yml | 39 ++-
flake.nix | 2 +
justfile | 12 +-
master/api.py | 29 +++
master/idempotency.py | 14 +-
shared/logger.py | 44 +++-
shared/openai.py | 20 ++
shared/protobufs/types/mlx/nn/__init__.pyi | 3 +
shared/pyproject.toml | 1 +
shared/types/api.py | 12 +
shared/types/common.py | 16 ++
shared/types/event_sourcing.py | 99 --------
shared/types/events/chunks.py | 93 +++++++
shared/types/events/common.py | 274 +++++++++++++++++++++
shared/types/events/events.py | 185 ++++++++++++++
shared/types/graphs/common.py | 171 +++++++++++++
shared/types/graphs/resource_graph.py | 17 ++
shared/types/models/common.py | 5 +
shared/types/models/metadata.py | 9 +
shared/types/models/model.py | 18 ++
shared/types/models/sources.py | 66 +++++
shared/types/networking/control_plane.py | 11 +
shared/types/networking/data_plane.py | 68 +++++
shared/types/networking/services.py | 29 +++
shared/types/networking/topology.py | 72 ++++++
shared/types/profiling/common.py | 54 ++++
shared/types/states/master.py | 85 +++++++
shared/types/states/shared.py | 30 +++
shared/types/states/worker.py | 17 ++
shared/types/tasks/common.py | 120 +++++++++
shared/types/worker/commands_runner.py | 102 ++++++++
shared/types/worker/common.py | 17 ++
shared/types/worker/downloads.py | 85 +++++++
shared/types/worker/instances.py | 35 +++
shared/types/worker/mlx.py | 13 +
shared/types/worker/resource_monitor.py | 73 ++++++
shared/types/worker/runners.py | 74 ++++++
shared/types/worker/shards.py | 54 ++++
shared/utils.py | 9 +
uv.lock | 148 +++++++++++
43 files changed, 2121 insertions(+), 134 deletions(-)
diff --git a/.github/actions/conditional-commit/conditional-commit.yml b/.github/actions/conditional-commit/action.yml
similarity index 100%
rename from .github/actions/conditional-commit/conditional-commit.yml
rename to .github/actions/conditional-commit/action.yml
diff --git a/.github/actions/lint-check/action.yml b/.github/actions/lint-check/action.yml
new file mode 100644
index 00000000..f666cae9
--- /dev/null
+++ b/.github/actions/lint-check/action.yml
@@ -0,0 +1,10 @@
+name: Lint Check
+
+description: "Check for lint errors"
+
+runs:
+ using: "composite"
+ steps:
+ - name: Lint check
+ run: nix develop -c just lint-check
+ shell: bash
diff --git a/.github/actions/verify-clean/action.yml b/.github/actions/verify-clean/action.yml
new file mode 100644
index 00000000..976e6a7d
--- /dev/null
+++ b/.github/actions/verify-clean/action.yml
@@ -0,0 +1,20 @@
+name: Verify Clean Working Tree
+
+description: "Fail the job if the previous step left the working tree dirty"
+
+inputs:
+ step:
+ description: "The name of the step that just executed"
+ required: true
+
+runs:
+ using: composite
+ steps:
+ - name: Check git diff
+ shell: bash
+ run: |
+ if ! git diff --quiet; then
+ echo "Error: ${{ inputs.step }} left working tree dirty." >&2
+ git --no-pager diff >&2
+ exit 1
+ fi
\ No newline at end of file
diff --git a/.github/workflows/pipeline.yml b/.github/workflows/pipeline.yml
index 6f3ba411..e2834848 100644
--- a/.github/workflows/pipeline.yml
+++ b/.github/workflows/pipeline.yml
@@ -14,12 +14,24 @@ jobs:
typecheck:
runs-on: ubuntu-22.04
steps:
+ - uses: actions/checkout@v4
+
+ - name: Configure git user
+ run: |
+ git config --local user.email "github-actions@users.noreply.github.com"
+ git config --local user.name "github-actions bot"
+ shell: bash
+
+ - uses: cachix/install-nix-action@v31
+ with:
+ github_access_token: ${{ secrets.GITHUB_TOKEN }}
+
- uses: ./.github/actions/typecheck
ci:
needs: typecheck
runs-on: ubuntu-22.04
permissions:
- contents: write
+ contents: read
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
steps:
@@ -38,27 +50,8 @@ jobs:
with:
github_access_token: ${{ secrets.GITHUB_TOKEN }}
- - uses: ./.github/actions/regenerate-protobufs
-
- - name: Commit regenerated protobufs
- uses: ./.github/actions/conditional-commit
- with:
- message: "chore(proto) regenerate protobufs"
-
- - uses: ./.github/actions/format
-
- - name: Commit formatted code
- uses: ./.github/actions/conditional-commit
+ - uses: ./.github/actions/verify-clean
with:
- message: "chore(format): format code"
-
- - uses: ./.github/actions/lint
+ step: regenerate-protobufs
- - name: Commit lint fixes
- uses: ./.github/actions/conditional-commit
- with:
- message: "chore(lint): fix linting errors"
-
- - name: Push changes
- run: git push
- shell: bash
+ - uses: ./.github/actions/lint-check
\ No newline at end of file
diff --git a/flake.nix b/flake.nix
index 2e1b6243..4ad5a219 100644
--- a/flake.nix
+++ b/flake.nix
@@ -22,6 +22,8 @@
pkgs.uv
pkgs.just
pkgs.protobuf
+ pkgs.rustc
+ pkgs.cargo
];
};
}
diff --git a/justfile b/justfile
index a2fe657a..fdffc979 100644
--- a/justfile
+++ b/justfile
@@ -1,6 +1,11 @@
regenerate-protobufs:
- protoc --proto_path=shared/protobufs/schemas --python_out=shared/protobufs/types --pyi_out=shared/protobufs/types shared/protobufs/schemas/*.proto
- uv run ruff format ./shared/protobufs/types
+ #!/usr/bin/env bash
+ if [ -f shared/protobufs/schemas/*.proto ]; then
+ protoc --proto_path=shared/protobufs/schemas --python_out=shared/protobufs/types --pyi_out=shared/protobufs/types shared/protobufs/schemas/*.proto
+ uv run ruff format ./shared/protobufs/types
+ else
+ echo "No .proto files found in shared/protobufs/schemas/"
+ fi
fmt:
uv run ruff format master worker shared engines/*
@@ -8,6 +13,9 @@ fmt:
lint:
uv run ruff check --fix master worker shared engines/*
+lint-check:
+ uv run ruff check master worker shared engines/*
+
test:
uv run pytest master worker shared engines/*
diff --git a/master/api.py b/master/api.py
new file mode 100644
index 00000000..50cc3bd3
--- /dev/null
+++ b/master/api.py
@@ -0,0 +1,29 @@
+from typing import Protocol
+
+from shared.types.models.common import ModelId
+from shared.types.models.model import ModelInfo
+from shared.types.models.sources import ModelSource
+from shared.types.networking.topology import ControlPlaneTopology, DataPlaneTopology
+from shared.types.worker.common import InstanceId
+from shared.types.worker.downloads import DownloadProgress
+from shared.types.worker.instances import Instance
+
+
+class ControlPlaneAPI(Protocol):
+ def get_control_plane_topology(self) -> ControlPlaneTopology: ...
+
+ def get_data_plane_topology(self) -> DataPlaneTopology: ...
+
+ def list_instances(self) -> list[Instance]: ...
+
+ def get_instance(self, instance_id: InstanceId) -> Instance: ...
+
+ def create_instance(self, model_id: ModelId) -> InstanceId: ...
+
+ def remove_instance(self, instance_id: InstanceId) -> None: ...
+
+ def get_model_data(self, model_id: ModelId) -> ModelInfo: ...
+
+ def download_model(self, model_id: ModelId, model_source: ModelSource) -> None: ...
+
+ def get_download_progress(self, model_id: ModelId) -> DownloadProgress: ...
diff --git a/master/idempotency.py b/master/idempotency.py
index 661d1e44..508cec6d 100644
--- a/master/idempotency.py
+++ b/master/idempotency.py
@@ -1,24 +1,28 @@
from hashlib import sha3_224 as hasher
from typing import Sequence, TypeVar
+from uuid import UUID
-from shared.types.event_sourcing import EventId, EventTypes, IdemKeyGenerator, State
+from shared.types.events.common import EventCategories, EventId, IdemKeyGenerator, State
-EventTypeT = TypeVar("EventTypeT", bound=EventTypes)
+EventCategoryT = TypeVar("EventCategoryT", bound=EventCategories)
-def get_idem_tag_generator(base: str) -> IdemKeyGenerator[EventTypeT]:
+def get_idem_tag_generator(base: str) -> IdemKeyGenerator[EventCategoryT]:
"""Generates idempotency keys for events.
The keys are generated by hashing the state sequence number against a base string.
You can pick any base string, **so long as it's not used in any other function that generates idempotency keys**.
"""
- def get_idem_keys(state: State[EventTypeT], num_keys: int) -> Sequence[EventId]:
+ def get_idem_keys(state: State[EventCategoryT], num_keys: int) -> Sequence[EventId]:
def recurse(n: int, last: bytes) -> Sequence[EventId]:
if n == 0:
return []
next_hash = hasher(last).digest()
- return (EventId(next_hash.hex()), *recurse(n - 1, next_hash))
+ return (
+ EventId(UUID(bytes=next_hash, version=4)),
+ *recurse(n - 1, next_hash),
+ )
initial_bytes = state.sequence_number.to_bytes(8, byteorder="big", signed=False)
return recurse(num_keys, initial_bytes)
diff --git a/shared/logger.py b/shared/logger.py
index 1d522fc2..659f551e 100644
--- a/shared/logger.py
+++ b/shared/logger.py
@@ -1,21 +1,48 @@
import logging
import logging.handlers
-from collections.abc import Sequence
+from collections.abc import Sequence, Set
+from enum import Enum
from queue import Queue
+from pydantic import BaseModel
from rich.logging import RichHandler
+class LogEntryType(str, Enum):
+ telemetry = "telemetry"
+ metrics = "metrics"
+ cluster = "cluster"
+
+
+class LogEntry(BaseModel):
+ event_type: Set[LogEntryType]
+
+
+class LogFilterByType(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()
+ LogEntry.model_validate_json(message)
+ return True
+
+
def configure_logger(
logger_name: str,
log_level: int = logging.INFO,
effect_handlers: Sequence[logging.Handler] | None = None,
) -> logging.Logger:
+ existing_logger = logging.Logger.manager.loggerDict.get(logger_name)
+ if existing_logger is not None:
+ raise RuntimeError(f"Logger with name '{logger_name}' already exists.")
+
logger = logging.getLogger(logger_name)
logger.setLevel(log_level)
logger.propagate = False
+ logging.raiseExceptions = True
- # If the named logger already has handlers, we assume it has been configured.
if logger.hasHandlers():
return logger
@@ -33,13 +60,20 @@ def configure_logger(
return logger
-def attach_to_queue(logger: logging.Logger, queue: Queue[logging.LogRecord]) -> None:
- logger.addHandler(logging.handlers.QueueHandler(queue))
+def attach_to_queue(
+ logger: logging.Logger,
+ filter_with: Sequence[logging.Filter],
+ queue: Queue[logging.LogRecord],
+) -> None:
+ handler = logging.handlers.QueueHandler(queue)
+ for log_filter in filter_with:
+ handler.addFilter(log_filter)
+ logger.addHandler(handler)
def create_queue_listener(
log_queue: Queue[logging.LogRecord],
- effect_handlers: list[logging.Handler],
+ effect_handlers: Sequence[logging.Handler],
) -> logging.handlers.QueueListener:
listener = logging.handlers.QueueListener(
log_queue, *effect_handlers, respect_handler_level=True
diff --git a/shared/openai.py b/shared/openai.py
new file mode 100644
index 00000000..0a0a546f
--- /dev/null
+++ b/shared/openai.py
@@ -0,0 +1,20 @@
+from typing import TYPE_CHECKING, Literal, TypeAlias, get_type_hints
+
+if TYPE_CHECKING:
+ import openai.types as openai_types
+ import openai.types.chat as openai_chat
+
+ types = openai_types
+ chat = openai_chat
+else:
+ types = None
+ chat = None
+
+FinishReason: TypeAlias = Literal[
+ "stop", "length", "tool_calls", "content_filter", "function_call"
+]
+assert (
+ get_type_hints(chat.chat_completion_chunk.Choice)["finish_reason"] == FinishReason
+), "Upstream changed Choice.finish_reason; update FinishReason alias."
+
+__all__ = ["types", "chat", "FinishReason"]
diff --git a/shared/protobufs/types/mlx/nn/__init__.pyi b/shared/protobufs/types/mlx/nn/__init__.pyi
new file mode 100644
index 00000000..464c4f1a
--- /dev/null
+++ b/shared/protobufs/types/mlx/nn/__init__.pyi
@@ -0,0 +1,3 @@
+from mlx.nn.layers import *
+from mlx.nn import init as init, losses as losses
+from mlx.nn.utils import average_gradients as average_gradients, value_and_grad as value_and_grad
\ No newline at end of file
diff --git a/shared/pyproject.toml b/shared/pyproject.toml
index c17f3dc7..d4ee919e 100644
--- a/shared/pyproject.toml
+++ b/shared/pyproject.toml
@@ -5,6 +5,7 @@ description = "Shared utilities for the Exo project"
readme = "README.md"
requires-python = ">=3.13"
dependencies = [
+ "openai>=1.93.0",
"pathlib>=1.0.1",
"protobuf>=6.31.1",
"pydantic>=2.11.7",
diff --git a/shared/types/api.py b/shared/types/api.py
new file mode 100644
index 00000000..f1bdefbf
--- /dev/null
+++ b/shared/types/api.py
@@ -0,0 +1,12 @@
+from typing import Literal
+
+from openai.types.chat.completion_create_params import CompletionCreateParams
+from pydantic import BaseModel
+
+from shared.types.tasks.common import TaskId
+
+
+class ChatTask(BaseModel):
+ task_id: TaskId
+ kind: Literal["chat"] = "chat"
+ task_data: CompletionCreateParams
diff --git a/shared/types/common.py b/shared/types/common.py
new file mode 100644
index 00000000..2c1b77ab
--- /dev/null
+++ b/shared/types/common.py
@@ -0,0 +1,16 @@
+from uuid import uuid4
+
+from pydantic import UUID4, Field
+from pydantic.dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class NewUUID:
+ uuid: UUID4 = Field(default_factory=lambda: uuid4())
+
+ def __hash__(self) -> int:
+ return hash(self.uuid)
+
+
+class NodeId(NewUUID):
+ pass
diff --git a/shared/types/event_sourcing.py b/shared/types/event_sourcing.py
deleted file mode 100644
index 33fc89e6..00000000
--- a/shared/types/event_sourcing.py
+++ /dev/null
@@ -1,99 +0,0 @@
-from typing import (
- Annotated,
- Callable,
- Generic,
- Literal,
- Protocol,
- Sequence,
- Tuple,
- TypeVar,
- get_args,
-)
-from uuid import UUID
-
-from pydantic import BaseModel, Field, TypeAdapter
-from pydantic.types import UuidVersion
-
-_EventId = Annotated[UUID, UuidVersion(4)]
-EventId = type("EventID", (UUID,), {})
-EventIdParser: TypeAdapter[EventId] = TypeAdapter(_EventId)
-
-EventTypes = Literal["create", "update", "delete"]
-EventTypeT = TypeVar("EventTypeT", bound=EventTypes)
-TEventType = TypeVar("TEventType", bound=EventTypes, covariant=True)
-
-
-class Event(BaseModel, Generic[TEventType]):
- event_type: TEventType
- idem_key: EventId
-
-
-class State(BaseModel, Generic[EventTypeT]):
- event_types: tuple[EventTypeT, ...] = get_args(EventTypeT)
- sequence_number: int = Field(default=0, ge=0)
-
-
-AnnotatedEventType = Annotated[EventTypes, Field(discriminator="event_type")]
-EventTypeParser: TypeAdapter[AnnotatedEventType] = TypeAdapter(AnnotatedEventType)
-
-Applicator = Callable[[State[EventTypeT], Event[TEventType]], State[EventTypeT]]
-Apply = Callable[[State[EventTypeT], Event[EventTypeT]], State[EventTypeT]]
-SagaApplicator = Callable[
- [State[EventTypeT], Event[TEventType]], Sequence[Event[EventTypeT]]
-]
-Saga = Callable[[State[EventTypeT], Event[EventTypeT]], Sequence[Event[EventTypeT]]]
-
-StateAndEvent = Tuple[State[EventTypeT], Event[EventTypeT]]
-EffectHandler = Callable[[StateAndEvent[EventTypeT], State[EventTypeT]], None]
-EventPublisher = Callable[[Event[EventTypeT]], None]
-
-
-class EventOutbox(Protocol):
- def send(self, events: Sequence[Event[EventTypeT]]) -> None: ...
-
-
-class EventProcessor(Protocol):
- def update(
- self,
- state: State[EventTypeT],
- apply: Apply[EventTypeT],
- effect_handlers: Sequence[EffectHandler[EventTypeT]],
- ) -> State[EventTypeT]: ...
-
-
-def get_saga_effect_handler(
- sagas: Saga[EventTypeT], event_publisher: EventPublisher[EventTypeT]
-) -> EffectHandler[EventTypeT]:
- def effect_handler(state_and_event: StateAndEvent[EventTypeT]) -> None:
- trigger_state, trigger_event = state_and_event
- for event in sagas(trigger_state, trigger_event):
- event_publisher(event)
-
- return lambda state_and_event, _: effect_handler(state_and_event)
-
-
-def get_effects_from_sagas(
- sagas: Sequence[Saga[EventTypeT]], event_publisher: EventPublisher[EventTypeT]
-) -> Sequence[EffectHandler[EventTypeT]]:
- return [get_saga_effect_handler(saga, event_publisher) for saga in sagas]
-
-
-IdemKeyGenerator = Callable[[State[EventTypeT], int], Sequence[EventId]]
-
-_CommandId = Annotated[UUID, UuidVersion(4)]
-CommandId = type("CommandID", (UUID,), {})
-CommandIdParser: TypeAdapter[CommandId] = TypeAdapter(_CommandId)
-
-CommandTypes = Literal["create", "update", "delete"]
-CommandTypeT = TypeVar("CommandTypeT", bound=EventTypes)
-TCommandType = TypeVar("TCommandType", bound=EventTypes, covariant=True)
-
-
-class Command(BaseModel, Generic[TEventType, TCommandType]):
- command_type: TCommandType
- idem_key: CommandId
-
-
-Decide = Callable[
- [State[EventTypeT], Command[TEventType, TCommandType]], Sequence[Event[EventTypeT]]
-]
diff --git a/shared/types/events/chunks.py b/shared/types/events/chunks.py
new file mode 100644
index 00000000..e75d6e1e
--- /dev/null
+++ b/shared/types/events/chunks.py
@@ -0,0 +1,93 @@
+from enum import Enum
+from typing import Annotated, Generic, Literal, TypeVar
+
+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]):
+ task_id: TaskId
+ idx: int
+ model: ModelId
+
+
+###
+
+
+class TokenChunkData(BaseModel):
+ text: str
+ token_id: int
+ finish_reason: FinishReason | None = None
+
+
+class ImageChunkData(BaseModel):
+ data: bytes
+
+
+###
+
+
+class TokenChunk(BaseChunk[ChunkType.token]):
+ chunk_data: TokenChunkData
+ chunk_type: Literal[ChunkType.token] = Field(default=ChunkType.token, frozen=True)
+
+
+class ImageChunk(BaseChunk[ChunkType.image]):
+ chunk_data: ImageChunkData
+ chunk_type: Literal[ChunkType.image] = Field(default=ChunkType.image, frozen=True)
+
+
+###
+
+GenerationChunk = Annotated[TokenChunk | ImageChunk, Field(discriminator="chunk_type")]
+GenerationChunkTypeAdapter: TypeAdapter[GenerationChunk] = TypeAdapter(GenerationChunk)
+
+# my_chunk: dict[str, Any] = TokenChunk(
+# task_id=TaskId('nicerid'),
+# idx=0,
+# chunk_data=TokenChunkData(
+# text='hello',
+# token_id=12,
+# ),
+# chunk_type=ChunkType.token,
+# model='llama-3.1',
+# ).model_dump()
+# print(my_chunk)
+# restored = GenerationChunkTypeAdapter.validate_python(my_chunk)
+# print(restored)
+
+#### OpenAI API Interfaces ###
+
+"""
+def send_task(task: Any) -> AsyncGenerator[GenerationChunk]:
+ # This is the 'command' - turns the task into an event and pushes to the event queue.
+ # Tokens are then read off the event queue and pushed back to the api via an AsyncGenerator.
+ ...
+
+def parse_chunk_to_openai_response(chunk: GenerationChunk) -> OpenAIResponse:
+ ...
+
+async def handle_task(task: Any) -> AsyncGenerator[OpenAIResponse]:
+ ## In our api call function, we will do:
+ generator: AsyncGenerator[GenerationChunk] = send_task(task)
+
+ async for chunk in generator:
+ yield parse_chunk_to_openai_response(chunk)
+"""
diff --git a/shared/types/events/common.py b/shared/types/events/common.py
new file mode 100644
index 00000000..6e5f78cf
--- /dev/null
+++ b/shared/types/events/common.py
@@ -0,0 +1,274 @@
+from enum import Enum, auto
+from typing import (
+ Annotated,
+ Callable,
+ Generic,
+ Protocol,
+ Sequence,
+ Tuple,
+ TypeVar,
+)
+
+from pydantic import BaseModel, Field, TypeAdapter, model_validator
+
+from shared.types.common import NewUUID, NodeId
+
+
+class EventId(NewUUID):
+ pass
+
+
+class TimerId(NewUUID):
+ pass
+
+
+class MLXEventTypes(str, Enum):
+ MLXInferenceSagaPrepare = "MLXInferenceSagaPrepare"
+ MLXInferenceSagaStartPrepare = "MLXInferenceSagaStartPrepare"
+
+
+class TaskEventTypes(str, Enum):
+ TaskCreated = "TaskCreated"
+ TaskUpdated = "TaskUpdated"
+ TaskDeleted = "TaskDeleted"
+
+
+class StreamingEventTypes(str, Enum):
+ ChunkGenerated = "ChunkGenerated"
+
+
+class InstanceEventTypes(str, Enum):
+ InstanceCreated = "InstanceCreated"
+ InstanceDeleted = "InstanceDeleted"
+ InstanceToBeReplacedAtomically = "InstanceToBeReplacedAtomically"
+ InstanceReplacedAtomically = "InstanceReplacedAtomically"
+ InstanceStatusUpdated = "InstanceStatusUpdated"
+
+
+class InstanceStateEventTypes(str, Enum):
+ InstanceRunnerStateUpdated = "InstanceRunnerStateUpdated"
+
+
+class NodePerformanceEventTypes(str, Enum):
+ NodePerformanceProfiled = "NodePerformanceProfiled"
+
+
+class DataPlaneEventTypes(str, Enum):
+ DataPlaneEdgeCreated = "DataPlaneEdgeCreated"
+ DataPlaneEdgeProfiled = "DataPlaneEdgeProfiled"
+ DataPlaneEdgeDeleted = "DataPlaneEdgeDeleted"
+
+
+class ControlPlaneEventTypes(str, Enum):
+ WorkerConnected = "WorkerConnected"
+ WorkerStatusUpdated = "WorkerStatusUpdated"
+ WorkerDisconnected = "WorkerDisconnected"
+
+
+class TimerEventTypes(str, Enum):
+ TimerCreated = "TimerCreated"
+ TimerFired = "TimerFired"
+
+
+class ResourceEventTypes(str, Enum):
+ ResourceProfiled = "ResourceProfiled"
+
+
+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,
+)
+
+
+class Event(BaseModel, Generic[PossibleEventOfEventTypeT]):
+ event_type: PossibleEventOfEventTypeT
+ event_category: EventCategories
+ event_id: EventId
+
+ def check_origin_id(self, origin_id: NodeId) -> bool:
+ return True
+
+
+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
+
+ @model_validator(mode="after")
+ def check_origin_id(self) -> "WrappedMessage[PossibleEventOfEventTypeT]":
+ if self.message.check_origin_id(self.origin_id):
+ 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)
+
+
+class State(BaseModel, Generic[CategoryOfEventsT_cov]):
+ event_category: CategoryOfEventsT_cov
+ sequence_number: int = Field(default=0, ge=0)
+
+
+AnnotatedEventType = Annotated[
+ Event[EventCategories], Field(discriminator="event_category")
+]
+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]],
+]
+Saga = Callable[
+ [State[CategoryOfEventsT_inv], Event[CategoryOfEventsT_inv]],
+ Sequence[Event[CategoryOfEventsT_inv]],
+]
+Apply = Callable[
+ [State[CategoryOfEventsT_inv], Event[CategoryOfEventsT_inv]],
+ State[CategoryOfEventsT_inv],
+]
+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):
+ 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):
+ def get_events_to_apply(
+ self, state: State[EventCategoryT]
+ ) -> Sequence[Event[EventCategoryT]]: ...
+
+
+def get_saga_effect_handler[EventCategoryT: EventCategories](
+ saga: Saga[EventCategoryT], event_publisher: EventPublisher[EventCategoryT]
+) -> EffectHandler[EventCategoryT]:
+ def effect_handler(state_and_event: StateAndEvent[EventCategoryT]) -> None:
+ trigger_state, trigger_event = state_and_event
+ for event in saga(trigger_state, trigger_event):
+ event_publisher(event)
+
+ return lambda state_and_event, _: effect_handler(state_and_event)
+
+
+def get_effects_from_sagas[EventCategoryT: EventCategories](
+ 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]]
+
+
+class CommandId(NewUUID):
+ pass
+
+
+class CommandTypes(str, Enum):
+ Create = "Create"
+ Update = "Update"
+ Delete = "Delete"
+
+
+class Command[EventCategoryT: EventCategories, 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]],
+]
diff --git a/shared/types/events/events.py b/shared/types/events/events.py
new file mode 100644
index 00000000..1f6422c8
--- /dev/null
+++ b/shared/types/events/events.py
@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+from typing import Any, Literal, Tuple
+
+from pydantic import BaseModel
+
+from shared.types.common import NodeId
+from shared.types.events.common import (
+ ControlPlaneEvent,
+ ControlPlaneEventTypes,
+ DataPlaneEvent,
+ DataPlaneEventTypes,
+ InstanceEvent,
+ InstanceEventTypes,
+ InstanceStateEvent,
+ InstanceStateEventTypes,
+ MLXEvent,
+ MLXEventTypes,
+ NodePerformanceEvent,
+ NodePerformanceEventTypes,
+ ResourceEvent,
+ ResourceEventTypes,
+ StreamingEvent,
+ StreamingEventTypes,
+ TaskEvent,
+ TaskEventTypes,
+ TimerEvent,
+ TimerEventTypes,
+ TimerId,
+)
+from shared.types.networking.control_plane import (
+ ControlPlaneEdgeId,
+ ControlPlaneEdgeType,
+)
+from shared.types.networking.data_plane import (
+ DataPlaneEdge,
+ DataPlaneEdgeId,
+ DataPlaneEdgeProfile,
+)
+from shared.types.profiling.common import NodePerformanceProfile, ProfiledResourceName
+from shared.types.tasks.common import (
+ TaskData,
+ TaskId,
+ TaskState,
+ TaskStatusIncompleteType,
+ 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
+
+
+class TimerData(BaseModel):
+ timer_id: TimerId
+
+
+class TaskCreated[TaskTypeT: TaskType](TaskEvent):
+ event_type: TaskEventTypes = TaskEventTypes.TaskCreated
+ task_id: TaskId
+ task_data: TaskData[TaskTypeT]
+ task_state: TaskState[Literal[TaskStatusIncompleteType.Pending], TaskTypeT]
+ 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
+ task_id: TaskId
+
+
+class InstanceCreated(InstanceEvent):
+ event_type: InstanceEventTypes = InstanceEventTypes.InstanceCreated
+ instance_id: InstanceId
+ instance_data: InstanceData
+ target_status: InstanceStatus
+
+
+class InstanceDeleted(InstanceEvent):
+ event_type: InstanceEventTypes = InstanceEventTypes.InstanceDeleted
+ instance_id: InstanceId
+
+
+class InstanceStatusUpdated(InstanceEvent):
+ event_type: InstanceEventTypes = InstanceEventTypes.InstanceStatusUpdated
+ instance_id: InstanceId
+ instance_status: InstanceStatus
+
+
+class InstanceRunnerStateUpdated(InstanceStateEvent):
+ event_type: InstanceStateEventTypes = (
+ InstanceStateEventTypes.InstanceRunnerStateUpdated
+ )
+ 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 MLXInferenceSagaPrepare(MLXEvent):
+ event_type: MLXEventTypes = MLXEventTypes.MLXInferenceSagaPrepare
+ task_id: TaskId
+ instance_id: InstanceId
+
+
+class MLXInferenceSagaStartPrepare(MLXEvent):
+ event_type: MLXEventTypes = MLXEventTypes.MLXInferenceSagaStartPrepare
+ task_id: TaskId
+ instance_id: InstanceId
+
+
+class NodePerformanceProfiled(NodePerformanceEvent):
+ event_type: NodePerformanceEventTypes = (
+ NodePerformanceEventTypes.NodePerformanceProfiled
+ )
+ node_id: NodeId
+ node_profile: NodePerformanceProfile
+
+
+class WorkerConnected(ControlPlaneEvent):
+ event_type: ControlPlaneEventTypes = ControlPlaneEventTypes.WorkerConnected
+ edge: DataPlaneEdge
+
+
+class WorkerStatusUpdated(ControlPlaneEvent):
+ event_type: ControlPlaneEventTypes = ControlPlaneEventTypes.WorkerStatusUpdated
+ node_id: NodeId
+ node_state: NodeStatus
+
+
+class WorkerDisconnected(ControlPlaneEvent):
+ event_type: ControlPlaneEventTypes = ControlPlaneEventTypes.WorkerConnected
+ vertex_id: ControlPlaneEdgeId
+
+
+class ChunkGenerated(StreamingEvent):
+ event_type: StreamingEventTypes = StreamingEventTypes.ChunkGenerated
+ task_id: TaskId
+ instance_id: InstanceId
+ chunk: Any
+
+
+class DataPlaneEdgeCreated(DataPlaneEvent):
+ event_type: DataPlaneEventTypes = DataPlaneEventTypes.DataPlaneEdgeCreated
+ vertex: ControlPlaneEdgeType
+
+
+class DataPlaneEdgeProfiled(DataPlaneEvent):
+ event_type: DataPlaneEventTypes = DataPlaneEventTypes.DataPlaneEdgeProfiled
+ edge_id: DataPlaneEdgeId
+ edge_profile: DataPlaneEdgeProfile
+
+
+class DataPlaneEdgeDeleted(DataPlaneEvent):
+ event_type: DataPlaneEventTypes = 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
diff --git a/shared/types/graphs/common.py b/shared/types/graphs/common.py
new file mode 100644
index 00000000..b43581fa
--- /dev/null
+++ b/shared/types/graphs/common.py
@@ -0,0 +1,171 @@
+from collections.abc import Mapping
+from typing import Callable, Generic, Protocol, Set, Tuple, TypeVar, overload
+
+from pydantic import BaseModel
+
+from shared.types.common import NewUUID
+
+EdgeTypeT = TypeVar("EdgeTypeT", covariant=True)
+VertexTypeT = TypeVar("VertexTypeT", covariant=True)
+EdgeIdT = TypeVar("EdgeIdT", bound=NewUUID)
+VertexIdT = TypeVar("VertexIdT", bound=NewUUID)
+
+
+class VertexData(BaseModel, Generic[VertexTypeT]):
+ vertex_type: VertexTypeT
+
+
+class EdgeData(BaseModel, Generic[EdgeTypeT]):
+ edge_type: EdgeTypeT
+
+
+class BaseEdge(BaseModel, Generic[EdgeTypeT, EdgeIdT, VertexIdT]):
+ edge_vertices: Tuple[VertexIdT, VertexIdT]
+ edge_data: EdgeData[EdgeTypeT]
+
+
+class BaseVertex(BaseModel, Generic[VertexTypeT, EdgeIdT]):
+ vertex_data: VertexData[VertexTypeT]
+
+
+class Vertex(
+ BaseVertex[VertexTypeT, EdgeIdT], Generic[VertexTypeT, EdgeIdT, VertexIdT]
+):
+ vertex_id: VertexIdT
+
+
+class Edge(
+ BaseEdge[EdgeTypeT, EdgeIdT, VertexIdT], Generic[EdgeTypeT, EdgeIdT, VertexIdT]
+):
+ edge_id: EdgeIdT
+
+
+class GraphData(BaseModel, Generic[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]):
+ edges: Mapping[EdgeIdT, EdgeData[EdgeTypeT]]
+ vertices: Mapping[VertexIdT, VertexData[VertexTypeT]]
+
+
+class GraphProtocol(Protocol, Generic[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]):
+ def list_edges(self) -> Set[EdgeIdT]: ...
+ def list_vertices(self) -> Set[VertexIdT]: ...
+ def get_vertices_from_edges(
+ self, edges: Set[EdgeIdT]
+ ) -> Mapping[EdgeIdT, Set[VertexIdT]]: ...
+ def get_edges_from_vertices(
+ self, vertices: Set[VertexIdT]
+ ) -> Mapping[VertexIdT, Set[EdgeIdT]]: ...
+ def get_edge_data(
+ self, edges: Set[EdgeIdT]
+ ) -> Mapping[EdgeIdT, EdgeData[EdgeTypeT]]: ...
+ def get_vertex_data(
+ self, vertices: Set[VertexIdT]
+ ) -> Mapping[VertexIdT, VertexData[VertexTypeT]]: ...
+
+
+class MutableGraphProtocol(GraphProtocol[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]):
+ def check_edges_exists(self, edge_id: EdgeIdT) -> bool: ...
+ def check_vertex_exists(self, vertex_id: VertexIdT) -> bool: ...
+ def _add_edge(self, edge_id: EdgeIdT, edge_data: EdgeData[EdgeTypeT]) -> None: ...
+ def _add_vertex(
+ self, vertex_id: VertexIdT, vertex_data: VertexData[VertexTypeT]
+ ) -> None: ...
+ def _remove_edge(self, edge_id: EdgeIdT) -> None: ...
+ def _remove_vertex(self, vertex_id: VertexIdT) -> None: ...
+ ###
+ @overload
+ def attach_edge(self, edge: Edge[EdgeTypeT, EdgeIdT, VertexIdT]) -> None: ...
+ @overload
+ def attach_edge(
+ self,
+ edge: Edge[EdgeTypeT, EdgeIdT, VertexIdT],
+ extra_vertex: Vertex[VertexTypeT, EdgeIdT, VertexIdT],
+ ) -> None: ...
+ def attach_edge(
+ self,
+ edge: Edge[EdgeTypeT, EdgeIdT, VertexIdT],
+ extra_vertex: Vertex[VertexTypeT, EdgeIdT, VertexIdT] | None = None,
+ ) -> None:
+ base_vertex = edge.edge_vertices[0]
+ target_vertex = edge.edge_vertices[1]
+ base_vertex_exists = self.check_vertex_exists(base_vertex)
+ target_vertex_exists = self.check_vertex_exists(target_vertex)
+
+ if not base_vertex_exists:
+ raise ValueError("Base Vertex Does Not Exist")
+
+ match (target_vertex_exists, extra_vertex is not None):
+ case (True, False):
+ raise ValueError("New Vertex Already Exists")
+ case (False, True):
+ if extra_vertex is None:
+ raise ValueError("BUG: Extra Vertex Must Be Provided")
+ self._add_vertex(extra_vertex.vertex_id, extra_vertex.vertex_data)
+ case (False, False):
+ raise ValueError(
+ "New Vertex Must Be Provided For Non-Existent Target Vertex"
+ )
+ case (True, True):
+ raise ValueError("New Vertex Already Exists")
+
+ self._add_edge(edge.edge_id, edge.edge_data)
+
+
+class Graph(
+ BaseModel,
+ Generic[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
+ GraphProtocol[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
+):
+ graph_data: GraphData[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]
+
+
+# the first element in the return value is the filtered graph; the second is the
+# (possibly empty) set of sub-graphs that were detached during filtering.
+def filter_by_edge_data(
+ graph: Graph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
+ keep: VertexIdT,
+ predicate: Callable[[EdgeData[EdgeTypeT]], bool],
+) -> Tuple[
+ Graph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
+ Set[Graph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]],
+]: ...
+
+
+# the first element in the return value is the filtered graph; the second is the
+# (possibly empty) set of sub-graphs that were detached during filtering.
+def filter_by_vertex_data(
+ graph: Graph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
+ keep: VertexIdT,
+ predicate: Callable[[VertexData[VertexTypeT]], bool],
+) -> Tuple[
+ Graph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
+ Set[Graph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]],
+]: ...
+
+
+def map_vertices_onto_graph(
+ vertices: Mapping[VertexIdT, VertexData[VertexTypeT]],
+ graph: Graph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
+) -> Tuple[Graph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT], Set[VertexIdT]]: ...
+
+
+def map_edges_onto_graph(
+ edges: Mapping[EdgeIdT, EdgeData[EdgeTypeT]],
+ graph: Graph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
+) -> Tuple[Graph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT], Set[EdgeIdT]]: ...
+
+
+def split_graph_by_edge(
+ graph: Graph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
+ edge: EdgeIdT,
+ keep: VertexIdT,
+) -> Tuple[
+ Graph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT],
+ Set[Graph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]],
+]: ...
+
+
+def merge_graphs_by_edge(
+ graphs: Set[Graph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT]],
+ edge: EdgeIdT,
+ keep: VertexIdT,
+) -> Tuple[Graph[EdgeTypeT, VertexTypeT, EdgeIdT, VertexIdT], Set[EdgeIdT]]: ...
diff --git a/shared/types/graphs/resource_graph.py b/shared/types/graphs/resource_graph.py
new file mode 100644
index 00000000..8f664507
--- /dev/null
+++ b/shared/types/graphs/resource_graph.py
@@ -0,0 +1,17 @@
+from collections.abc import Mapping
+
+from pydantic import BaseModel
+
+from shared.types.common import NodeId
+from shared.types.networking.topology import ControlPlaneTopology, DataPlaneTopology
+from shared.types.profiling.common import NodePerformanceProfile
+
+
+class ResourceGraph(BaseModel): ...
+
+
+def get_graph_of_compute_resources(
+ control_plane_topology: ControlPlaneTopology,
+ data_plane_topology: DataPlaneTopology,
+ node_profiles: Mapping[NodeId, NodePerformanceProfile],
+) -> ResourceGraph: ...
diff --git a/shared/types/models/common.py b/shared/types/models/common.py
new file mode 100644
index 00000000..05e82a34
--- /dev/null
+++ b/shared/types/models/common.py
@@ -0,0 +1,5 @@
+from shared.types.common import NewUUID
+
+
+class ModelId(NewUUID):
+ pass
diff --git a/shared/types/models/metadata.py b/shared/types/models/metadata.py
new file mode 100644
index 00000000..1d42d3dc
--- /dev/null
+++ b/shared/types/models/metadata.py
@@ -0,0 +1,9 @@
+from typing import Annotated
+
+from pydantic import BaseModel, PositiveInt
+
+
+class ModelMetadata(BaseModel):
+ pretty_name: str
+ storage_size_kilobytes: Annotated[int, PositiveInt]
+ n_layers: Annotated[int, PositiveInt]
diff --git a/shared/types/models/model.py b/shared/types/models/model.py
new file mode 100644
index 00000000..faa7c3ad
--- /dev/null
+++ b/shared/types/models/model.py
@@ -0,0 +1,18 @@
+from typing import Sequence, final
+
+from pydantic import BaseModel, TypeAdapter
+
+from shared.types.models.common import ModelId
+from shared.types.models.metadata import ModelMetadata
+from shared.types.models.sources import ModelSource
+
+
+@final
+# Concerned by the naming here; model could also be an instance of a model.
+class ModelInfo(BaseModel):
+ model_id: ModelId
+ model_sources: Sequence[ModelSource]
+ model_metadata: ModelMetadata
+
+
+ModelIdAdapter: TypeAdapter[ModelId] = TypeAdapter(ModelId)
diff --git a/shared/types/models/sources.py b/shared/types/models/sources.py
new file mode 100644
index 00000000..8f636a26
--- /dev/null
+++ b/shared/types/models/sources.py
@@ -0,0 +1,66 @@
+from enum import Enum
+from typing import Annotated, Any, Generic, Literal, TypeVar, Union, final
+
+from pydantic import AnyHttpUrl, BaseModel, Field, TypeAdapter
+
+from shared.types.models.common import ModelId
+
+
+class SourceType(str, Enum):
+ HuggingFace = "HuggingFace"
+ GitHub = "GitHub"
+
+
+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]):
+ model_uuid: ModelId
+ source_type: T
+ source_format: S
+ source_data: Any
+
+
+@final
+class HuggingFaceModelSourceData(BaseModel):
+ path: RepoPath
+
+
+@final
+class GitHubModelSourceData(BaseModel):
+ url: AnyHttpUrl
+
+
+@final
+class HuggingFaceModelSource(
+ BaseModelSource[SourceType.HuggingFace, SourceFormatType.HuggingFaceTransformers]
+):
+ source_type: Literal[SourceType.HuggingFace] = SourceType.HuggingFace
+ source_format: Literal[SourceFormatType.HuggingFaceTransformers] = (
+ SourceFormatType.HuggingFaceTransformers
+ )
+ source_data: HuggingFaceModelSourceData
+
+
+@final
+class GitHubModelSource(BaseModelSource[SourceType.GitHub, S]):
+ source_type: Literal[SourceType.GitHub] = SourceType.GitHub
+ source_data: GitHubModelSourceData
+
+
+_ModelSource = Annotated[
+ Union[
+ HuggingFaceModelSource,
+ GitHubModelSource[SourceFormatType.HuggingFaceTransformers],
+ ],
+ Field(discriminator="source_type"),
+]
+ModelSource = BaseModelSource[SourceType, SourceFormatType]
+ModelSourceAdapter: TypeAdapter[ModelSource] = TypeAdapter(_ModelSource)
diff --git a/shared/types/networking/control_plane.py b/shared/types/networking/control_plane.py
new file mode 100644
index 00000000..574ff097
--- /dev/null
+++ b/shared/types/networking/control_plane.py
@@ -0,0 +1,11 @@
+from typing import TypeAlias
+
+from shared.types.common import NewUUID, NodeId
+from shared.types.graphs.common import Edge
+
+
+class ControlPlaneEdgeId(NewUUID):
+ pass
+
+
+ControlPlaneEdgeType: TypeAlias = Edge[None, ControlPlaneEdgeId, NodeId]
diff --git a/shared/types/networking/data_plane.py b/shared/types/networking/data_plane.py
new file mode 100644
index 00000000..9c570973
--- /dev/null
+++ b/shared/types/networking/data_plane.py
@@ -0,0 +1,68 @@
+from enum import Enum
+from typing import Annotated, Literal, TypeVar, Union, final
+
+from pydantic import BaseModel, Field, IPvAnyAddress, TypeAdapter
+
+from shared.types.common import NewUUID, NodeId
+from shared.types.graphs.common import Edge
+
+
+class DataPlaneEdgeId(NewUUID):
+ pass
+
+
+class AddressingProtocol(str, Enum):
+ IPvAnyAddress = "IPvAnyAddress"
+
+
+class ApplicationProtocol(str, Enum):
+ MLX = "MLX"
+
+
+AdP = TypeVar("AdP", bound=AddressingProtocol)
+ApP = TypeVar("ApP", bound=ApplicationProtocol)
+
+
+@final
+class DataPlaneEdgeProfile(BaseModel):
+ throughput: float
+ latency: float
+ jitter: float
+
+
+class CommonDataPlaneEdgeData(BaseModel):
+ edge_data_transfer_rate: DataPlaneEdgeProfile | None = None
+
+
+class MlxEdgeMetadata(BaseModel):
+ source_ip: IPvAnyAddress
+ sink_ip: IPvAnyAddress
+
+
+class BaseDataPlaneEdgeData[AdP: AddressingProtocol, ApP: ApplicationProtocol](
+ BaseModel
+):
+ addressing_protocol: AdP
+ application_protocol: ApP
+ common_data: CommonDataPlaneEdgeData
+
+
+class MlxEdge(
+ BaseDataPlaneEdgeData[AddressingProtocol.IPvAnyAddress, ApplicationProtocol.MLX]
+):
+ addressing_protocol: Literal[AddressingProtocol.IPvAnyAddress] = (
+ AddressingProtocol.IPvAnyAddress
+ )
+ application_protocol: Literal[ApplicationProtocol.MLX] = ApplicationProtocol.MLX
+ mlx_metadata: MlxEdgeMetadata
+
+
+DataPlaneEdgeData = Union[MlxEdge]
+
+_DataPlaneEdgeData = Annotated[
+ DataPlaneEdgeData,
+ Field(discriminator="addressing_protocol"),
+]
+DataPlaneEdgeAdapter: TypeAdapter[DataPlaneEdgeData] = TypeAdapter(_DataPlaneEdgeData)
+
+DataPlaneEdge = Edge[DataPlaneEdgeData, DataPlaneEdgeId, NodeId]
diff --git a/shared/types/networking/services.py b/shared/types/networking/services.py
new file mode 100644
index 00000000..01655d15
--- /dev/null
+++ b/shared/types/networking/services.py
@@ -0,0 +1,29 @@
+from typing import Callable, NewType, Protocol
+
+from shared.types.networking.control_plane import (
+ ControlPlaneEdgeId,
+ ControlPlaneEdgeType,
+)
+
+TopicName = NewType("TopicName", str)
+
+PubSubMessageHandler = Callable[[TopicName, object], None]
+NodeConnectedHandler = Callable[
+ [
+ ControlPlaneEdgeId,
+ ControlPlaneEdgeType,
+ ],
+ None,
+]
+NodeDisconnectedHandler = Callable[[ControlPlaneEdgeId], None]
+
+
+class DiscoveryService(Protocol):
+ def on_node_connected(self, handler: NodeConnectedHandler) -> None: ...
+ def on_node_disconnected(self, handler: NodeDisconnectedHandler) -> None: ...
+
+
+class PubSubService(Protocol):
+ def on_message_received(
+ self, topic_name: TopicName, handler: PubSubMessageHandler
+ ) -> None: ...
diff --git a/shared/types/networking/topology.py b/shared/types/networking/topology.py
new file mode 100644
index 00000000..61e8900b
--- /dev/null
+++ b/shared/types/networking/topology.py
@@ -0,0 +1,72 @@
+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,
+ DataPlaneEdgeId,
+)
+from shared.types.worker.common import NodeStatus
+
+
+class DataPlaneTopology(
+ Graph[
+ DataPlaneEdgeData,
+ None,
+ DataPlaneEdgeId,
+ NodeId,
+ ]
+):
+ graph_data: GraphData[
+ DataPlaneEdgeData,
+ None,
+ DataPlaneEdgeId,
+ NodeId,
+ ]
+
+
+class OrphanedPartOfDataPlaneTopology(
+ Graph[
+ DataPlaneEdgeData,
+ None,
+ DataPlaneEdgeId,
+ NodeId,
+ ]
+):
+ graph_data: GraphData[
+ DataPlaneEdgeData,
+ None,
+ DataPlaneEdgeId,
+ NodeId,
+ ]
+
+
+class ControlPlaneTopology(
+ Graph[
+ None,
+ NodeStatus,
+ ControlPlaneEdgeId,
+ NodeId,
+ ]
+):
+ graph_data: GraphData[
+ None,
+ NodeStatus,
+ ControlPlaneEdgeId,
+ NodeId,
+ ]
+
+
+class OrphanedPartOfControlPlaneTopology(
+ Graph[
+ None,
+ NodeStatus,
+ ControlPlaneEdgeId,
+ NodeId,
+ ]
+):
+ graph_data: GraphData[
+ None,
+ NodeStatus,
+ ControlPlaneEdgeId,
+ NodeId,
+ ]
diff --git a/shared/types/profiling/common.py b/shared/types/profiling/common.py
new file mode 100644
index 00000000..1b318cc7
--- /dev/null
+++ b/shared/types/profiling/common.py
@@ -0,0 +1,54 @@
+from enum import Enum
+from typing import Annotated, Generic, Literal, TypeVar
+
+from pydantic import BaseModel, Field, TypeAdapter
+
+
+class ProfiledResourceName(str, Enum):
+ memory = "memory"
+ system = "system"
+
+
+ProfiledResourceT = TypeVar(name="ProfiledResourceT", bound=ProfiledResourceName)
+
+
+class BasePerformanceProfile(BaseModel, Generic[ProfiledResourceT]):
+ """
+ Details a single resource (or resource type) that is being monitored by the resource monitor.
+ """
+
+
+class MemoryPerformanceProfile(BasePerformanceProfile[ProfiledResourceName.memory]):
+ resource_name: Literal[ProfiledResourceName.memory] = Field(
+ default=ProfiledResourceName.memory, frozen=True
+ )
+ ram_total: int
+ ram_used: int
+ swap_total: int
+ swap_used: int
+
+
+class NetworkInterfaceInfo(BaseModel):
+ name: str
+ ip_address: str
+ type: str
+
+
+class SystemPerformanceProfile(BasePerformanceProfile[ProfiledResourceName.system]):
+ resource_name: Literal[ProfiledResourceName.system] = Field(
+ default=ProfiledResourceName.system, frozen=True
+ )
+ model_id: str
+ chip_id: str
+ memory: int
+ network_interfaces: list[NetworkInterfaceInfo] = Field(default_factory=list)
+
+
+NodePerformanceProfile = Annotated[
+ MemoryPerformanceProfile | SystemPerformanceProfile,
+ Field(discriminator="resource_name"),
+]
+
+NodePerformanceProfileTypeAdapter: TypeAdapter[NodePerformanceProfile] = TypeAdapter(
+ NodePerformanceProfile
+)
diff --git a/shared/types/states/master.py b/shared/types/states/master.py
new file mode 100644
index 00000000..e1233b11
--- /dev/null
+++ b/shared/types/states/master.py
@@ -0,0 +1,85 @@
+from collections.abc import Mapping, Sequence
+from enum import Enum
+from queue import Queue
+from typing import Generic, TypeVar
+
+from pydantic import BaseModel
+
+from shared.types.common import NodeId
+from shared.types.events.common import (
+ Event,
+ EventCategories,
+ State,
+)
+from shared.types.graphs.resource_graph import ResourceGraph
+from shared.types.networking.data_plane import (
+ DataPlaneEdge,
+ DataPlaneEdgeId,
+)
+from shared.types.networking.topology import (
+ ControlPlaneTopology,
+ DataPlaneTopology,
+ OrphanedPartOfControlPlaneTopology,
+ OrphanedPartOfDataPlaneTopology,
+)
+from shared.types.profiling.common import NodePerformanceProfile
+from shared.types.states.shared import SharedState
+from shared.types.tasks.common import TaskData, TaskType
+from shared.types.worker.instances import InstanceData, InstanceId
+
+
+class ExternalCommand(BaseModel): ...
+
+
+class CachePolicyType(str, Enum):
+ KeepAll = "KeepAll"
+
+
+CachePolicyTypeT = TypeVar("CachePolicyTypeT", bound=CachePolicyType)
+
+
+class CachePolicy(BaseModel, Generic[CachePolicyTypeT]):
+ policy_type: CachePolicyTypeT
+
+
+class NodePerformanceProfileState(State[EventCategories.NodePerformanceEventTypes]):
+ node_profiles: Mapping[NodeId, NodePerformanceProfile]
+
+
+class DataPlaneNetworkState(State[EventCategories.DataPlaneEventTypes]):
+ topology: DataPlaneTopology
+ 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]
+
+ 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[TaskType]]
+ job_outbox: Queue[TaskData[TaskType]]
+ cache_policy: CachePolicy[CachePolicyType]
+
+
+def get_shard_assignments(
+ inbox: Queue[ExternalCommand],
+ outbox: Queue[ExternalCommand],
+ resource_graph: ResourceGraph,
+ current_instances: Mapping[InstanceId, InstanceData],
+ cache_policy: CachePolicy[CachePolicyType],
+) -> Mapping[InstanceId, InstanceData]: ...
+
+
+def get_transition_events(
+ current_instances: Mapping[InstanceId, InstanceData],
+ target_instances: Mapping[InstanceId, InstanceData],
+) -> Sequence[Event[EventCategories]]: ...
diff --git a/shared/types/states/shared.py b/shared/types/states/shared.py
new file mode 100644
index 00000000..75e3140e
--- /dev/null
+++ b/shared/types/states/shared.py
@@ -0,0 +1,30 @@
+from collections.abc import Mapping
+from typing import 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.worker.common import InstanceId
+from shared.types.worker.instances import BaseInstance
+
+
+class KnownInstances(State[EventCategories.InstanceStateEventTypes]):
+ instances: Mapping[InstanceId, BaseInstance]
+
+
+class Tasks(State[EventCategories.TaskEventTypes]):
+ tasks: Mapping[TaskId, Task[TaskType, TaskStatusType]]
+
+
+class SharedState(BaseModel):
+ node_id: NodeId
+ known_instances: KnownInstances
+ compute_tasks: Tasks
+
+ def get_node_id(self) -> NodeId: ...
+
+ def get_tasks_by_instance(
+ self, instance_id: InstanceId
+ ) -> Sequence[Task[TaskType, TaskStatusType]]: ...
diff --git a/shared/types/states/worker.py b/shared/types/states/worker.py
new file mode 100644
index 00000000..699ecb84
--- /dev/null
+++ b/shared/types/states/worker.py
@@ -0,0 +1,17 @@
+from collections.abc import Mapping
+
+from shared.types.common import NodeId
+from shared.types.events.common import (
+ EventCategories,
+ State,
+)
+from shared.types.states.shared import SharedState
+from shared.types.worker.common import NodeStatus
+
+
+class NodeStatusState(State[EventCategories.ControlPlaneEventTypes]):
+ node_status: Mapping[NodeId, NodeStatus]
+
+
+class WorkerState(SharedState):
+ node_status: NodeStatusState
diff --git a/shared/types/tasks/common.py b/shared/types/tasks/common.py
new file mode 100644
index 00000000..7e58c35f
--- /dev/null
+++ b/shared/types/tasks/common.py
@@ -0,0 +1,120 @@
+from collections.abc import Mapping
+from enum import Enum
+from typing import Annotated, Generic, Literal, TypeVar, Union
+
+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
+
+
+class TaskId(NewUUID):
+ pass
+
+
+class TaskType(str, Enum):
+ ChatCompletionNonStreaming = "ChatCompletionNonStreaming"
+ ChatCompletionStreaming = "ChatCompletionStreaming"
+
+
+TaskTypeT = TypeVar("TaskTypeT", bound=TaskType, covariant=True)
+
+
+class TaskData(BaseModel, Generic[TaskTypeT]): ...
+
+
+class ChatCompletionNonStreamingTask(TaskData[TaskType.ChatCompletionNonStreaming]):
+ task_type: Literal[TaskType.ChatCompletionNonStreaming] = (
+ TaskType.ChatCompletionNonStreaming
+ )
+ task_data: openai.completion_create_params.CompletionCreateParams
+
+
+class ChatCompletionStreamingTask(TaskData[TaskType.ChatCompletionStreaming]):
+ task_type: Literal[TaskType.ChatCompletionStreaming] = (
+ TaskType.ChatCompletionStreaming
+ )
+ task_data: openai.completion_create_params.CompletionCreateParams
+
+
+class TaskStatusIncompleteType(str, Enum):
+ Pending = "Pending"
+ Running = "Running"
+ Failed = "Failed"
+
+
+class TaskStatusCompleteType(str, Enum):
+ Complete = "Complete"
+
+
+TaskStatusType = Union[TaskStatusIncompleteType, TaskStatusCompleteType]
+
+
+class TaskArtifact[TaskTypeT: TaskType, TaskStatusTypeT: TaskStatusType](BaseModel): ...
+
+
+class IncompleteTaskArtifact[TaskTypeT: TaskType](
+ TaskArtifact[TaskTypeT, TaskStatusIncompleteType]
+):
+ 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]
+
+
+class TaskState[TaskStatusTypeT: TaskStatusType, TaskTypeT: TaskType](BaseModel):
+ task_status: TaskStatusUpdate[TaskStatusTypeT]
+ task_artifact: TaskArtifact[TaskTypeT, TaskStatusTypeT]
+
+
+class BaseTask[TaskTypeT: TaskType, TaskStatusTypeT: TaskStatusType](BaseModel):
+ task_type: TaskTypeT
+ task_data: TaskData[TaskTypeT]
+ task_state: TaskState[TaskStatusTypeT, TaskTypeT]
+ on_instance: InstanceId
+
+
+BaseTaskAnnotated = Annotated[
+ Union[
+ BaseTask[Literal[TaskType.ChatCompletionNonStreaming], TaskStatusType],
+ BaseTask[Literal[TaskType.ChatCompletionStreaming], TaskStatusType],
+ ],
+ Field(discriminator="task_type"),
+]
+
+BaseTaskValidator: TypeAdapter[BaseTask[TaskType, TaskStatusType]] = TypeAdapter(
+ BaseTaskAnnotated
+)
+
+
+class Task[TaskTypeT: TaskType, TaskStatusTypeT: TaskStatusType](
+ BaseTask[TaskTypeT, TaskStatusTypeT]
+):
+ task_id: TaskId
diff --git a/shared/types/worker/commands_runner.py b/shared/types/worker/commands_runner.py
new file mode 100644
index 00000000..7f636588
--- /dev/null
+++ b/shared/types/worker/commands_runner.py
@@ -0,0 +1,102 @@
+from enum import Enum
+from typing import Annotated, Generic, Literal, TypeVar
+
+from pydantic import BaseModel, Field, TypeAdapter
+
+from shared.openai import FinishReason
+from shared.types.api import ChatTask
+from shared.types.worker.mlx import Host
+from shared.types.worker.shards import PartitionStrategy, ShardMetadata
+
+## Messages passed TO the runner
+
+
+class MessageType(str, Enum):
+ Setup = "setup"
+ ChatTask = "chat_task"
+ Exit = "exit"
+
+
+MT = TypeVar(name="MT", bound=MessageType)
+
+
+class BaseRunnerMessage(BaseModel, Generic[MT]):
+ pass
+
+
+class SetupMessage(BaseRunnerMessage[MessageType.Setup]):
+ type: Literal[MessageType.Setup] = Field(default=MessageType.Setup, frozen=True)
+ model_shard_meta: ShardMetadata[PartitionStrategy]
+ hosts: list[Host]
+
+
+class ChatTaskMessage(BaseRunnerMessage[MessageType.ChatTask]):
+ type: Literal[MessageType.ChatTask] = Field(
+ default=MessageType.ChatTask, frozen=True
+ )
+ task: ChatTask
+
+
+class ExitMessage(BaseRunnerMessage[MessageType.Exit]):
+ type: Literal[MessageType.Exit] = Field(default=MessageType.Exit, frozen=True)
+
+
+RunnerMessage = Annotated[
+ SetupMessage | ChatTaskMessage | ExitMessage, Field(discriminator="type")
+]
+RunnerMessageTypeAdapter: TypeAdapter[RunnerMessage] = TypeAdapter(RunnerMessage)
+
+## Responses passed FROM the runner
+
+
+class RunnerResponseType(str, Enum):
+ GenerationResponse = "generation_response"
+ FinishedResponse = "finished_response"
+ PrintResponse = "print_response"
+ ErrorResponse = "error_response"
+
+
+RRT = TypeVar(name="RRT", bound=RunnerResponseType)
+
+
+class BaseRunnerResponse(BaseModel, Generic[RRT]):
+ pass
+
+
+class GenerationResponse(BaseRunnerResponse[RunnerResponseType.GenerationResponse]):
+ type: Literal[RunnerResponseType.GenerationResponse] = Field(
+ default=RunnerResponseType.GenerationResponse, frozen=True
+ )
+ text: str
+ token: int
+ # logprobs: Optional[list[float]] = None # too big. we can change to be top-k
+ finish_reason: FinishReason | None = None
+
+
+class PrintResponse(BaseRunnerResponse[RunnerResponseType.PrintResponse]):
+ type: Literal[RunnerResponseType.PrintResponse] = Field(
+ default=RunnerResponseType.PrintResponse, frozen=True
+ )
+ text: str
+
+
+class FinishedResponse(BaseRunnerResponse[RunnerResponseType.FinishedResponse]):
+ type: Literal[RunnerResponseType.FinishedResponse] = Field(
+ default=RunnerResponseType.FinishedResponse, frozen=True
+ )
+
+
+class ErrorResponse(BaseRunnerResponse[RunnerResponseType.ErrorResponse]):
+ type: Literal[RunnerResponseType.ErrorResponse] = Field(
+ default=RunnerResponseType.ErrorResponse, frozen=True
+ )
+ error_type: str
+ error_message: str
+ traceback: str | None = None
+
+
+RunnerResponse = Annotated[
+ GenerationResponse | PrintResponse | FinishedResponse | ErrorResponse,
+ Field(discriminator="type"),
+]
+RunnerResponseTypeAdapter: TypeAdapter[RunnerResponse] = TypeAdapter(RunnerResponse)
diff --git a/shared/types/worker/common.py b/shared/types/worker/common.py
new file mode 100644
index 00000000..5fa78f74
--- /dev/null
+++ b/shared/types/worker/common.py
@@ -0,0 +1,17 @@
+from enum import Enum
+
+from shared.types.common import NewUUID
+
+
+class InstanceId(NewUUID):
+ pass
+
+
+class RunnerId(NewUUID):
+ pass
+
+
+class NodeStatus(str, Enum):
+ Idle = "Idle"
+ Running = "Running"
+ Paused = "Paused"
diff --git a/shared/types/worker/downloads.py b/shared/types/worker/downloads.py
new file mode 100644
index 00000000..c88b2d57
--- /dev/null
+++ b/shared/types/worker/downloads.py
@@ -0,0 +1,85 @@
+from enum import Enum
+from typing import (
+ Annotated,
+ Callable,
+ Generic,
+ Literal,
+ NewType,
+ Sequence,
+ TypeVar,
+ Union,
+)
+
+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 PartitionStrategy, ShardMetadata
+
+
+class DownloadProgressData(BaseModel):
+ total_bytes: Annotated[int, PositiveInt]
+ downloaded_bytes: Annotated[int, PositiveInt]
+
+
+class DownloadStatus(str, Enum):
+ Pending = "Pending"
+ Downloading = "Downloading"
+ Completed = "Completed"
+ Failed = "Failed"
+
+
+DownloadStatusT = TypeVar("DownloadStatusT", bound=DownloadStatus)
+
+
+class BaseDownloadProgress(BaseModel, Generic[DownloadStatusT]):
+ node_id: NodeId
+ download_status: DownloadStatusT
+
+
+class DownloadPending(BaseDownloadProgress[DownloadStatus.Pending]):
+ download_status: Literal[DownloadStatus.Pending] = Field(DownloadStatus.Pending)
+
+
+class DownloadCompleted(BaseDownloadProgress[DownloadStatus.Completed]):
+ download_status: Literal[DownloadStatus.Completed] = Field(DownloadStatus.Completed)
+
+
+class DownloadFailed(BaseDownloadProgress[DownloadStatus.Failed]):
+ download_status: Literal[DownloadStatus.Failed] = Field(DownloadStatus.Failed)
+ error_message: str
+
+
+class DownloadOngoing(BaseDownloadProgress[DownloadStatus.Downloading]):
+ download_status: Literal[DownloadStatus.Downloading] = Field(
+ DownloadStatus.Downloading
+ )
+ download_progress: DownloadProgressData
+
+
+DownloadProgress = Annotated[
+ Union[
+ DownloadPending,
+ DownloadCompleted,
+ DownloadFailed,
+ DownloadOngoing,
+ ],
+ Field(discriminator="download_status"),
+]
+
+
+BytesToDownload = NewType("BytesToDownload", int)
+BytesDownloaded = NewType("BytesDownloaded", int)
+
+DownloadEffectHandler = Callable[
+ [ModelId, DownloadStatus, BytesToDownload, BytesDownloaded], None
+]
+
+
+def download_shard(
+ model_id: ModelId,
+ model_source: ModelSource,
+ shard_meta: ShardMetadata[PartitionStrategy],
+ effect_handlers: Sequence[DownloadEffectHandler],
+) -> None: ...
diff --git a/shared/types/worker/instances.py b/shared/types/worker/instances.py
new file mode 100644
index 00000000..f23b5807
--- /dev/null
+++ b/shared/types/worker/instances.py
@@ -0,0 +1,35 @@
+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):
+ ACTIVE = "active"
+ INACTIVE = "inactive"
+
+
+class InstanceState(BaseModel):
+ runner_states: Mapping[RunnerId, RunnerState[RunnerStateType]]
+
+
+class InstanceData(BaseModel):
+ shard_assignments: ShardAssignments
+
+
+class BaseInstance(BaseModel):
+ instance_data: InstanceData
+ instance_state: InstanceState
+ instance_status: InstanceStatus
+
+
+class Instance(BaseInstance):
+ instance_id: InstanceId
diff --git a/shared/types/worker/mlx.py b/shared/types/worker/mlx.py
new file mode 100644
index 00000000..496ef369
--- /dev/null
+++ b/shared/types/worker/mlx.py
@@ -0,0 +1,13 @@
+from pydantic import BaseModel, field_validator
+
+
+# TODO: Is this the right place for this? Host is consumed by worker, but typically stored in the master
+class Host(BaseModel):
+ host: str
+ port: int
+
+ @field_validator("port")
+ def check_port(self, v: int) -> int:
+ if not (0 <= v <= 65535):
+ raise ValueError("Port must be between 0 and 65535")
+ return v
diff --git a/shared/types/worker/resource_monitor.py b/shared/types/worker/resource_monitor.py
new file mode 100644
index 00000000..96eba8d2
--- /dev/null
+++ b/shared/types/worker/resource_monitor.py
@@ -0,0 +1,73 @@
+import asyncio
+from abc import ABC, abstractmethod
+from collections.abc import Coroutine
+from typing import Callable, Set
+
+from shared.types.events.events import ResourceProfiled
+from shared.types.profiling.common import (
+ MemoryPerformanceProfile,
+ NodePerformanceProfile,
+ SystemPerformanceProfile,
+)
+
+
+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
+
+ @abstractmethod
+ async def collect(self) -> NodePerformanceProfile: ...
+
+
+class SystemResourceCollector(ResourceCollector):
+ def __init__(self):
+ super().__init__("system")
+
+ @abstractmethod
+ async def collect(self) -> SystemPerformanceProfile: ...
+
+
+class MemoryResourceCollector(ResourceCollector):
+ def __init__(self):
+ super().__init__("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
+
+ # 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
+ ]
+ return await asyncio.gather(*tasks)
+
+ async def collect(self) -> None:
+ profiles = await self._collect()
+ for profile in profiles:
+ for effect_handler in self.effect_handlers:
+ effect_handler(profile)
diff --git a/shared/types/worker/runners.py b/shared/types/worker/runners.py
new file mode 100644
index 00000000..c7528094
--- /dev/null
+++ b/shared/types/worker/runners.py
@@ -0,0 +1,74 @@
+from collections.abc import Mapping, Sequence
+from enum import Enum
+from typing import Generic, Literal, TypeVar
+
+from pydantic import BaseModel, 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 PartitionStrategy, ShardMetadata
+
+
+class RunnerStateType(str, Enum):
+ Rejected = "Rejected"
+ Starting = "Starting"
+ Downloading = "Downloading"
+ Running = "Running"
+ Failed = "Failed"
+
+
+RunnerStateTypeT = TypeVar("RunnerStateTypeT", bound=RunnerStateType)
+
+
+class RunnerState(BaseModel, Generic[RunnerStateTypeT]):
+ runner_state: RunnerStateTypeT
+
+
+class RejectedRunnerState(RunnerState[RunnerStateType.Rejected]):
+ runner_state: Literal[RunnerStateType.Rejected]
+
+
+class StartingRunnerState(RunnerState[RunnerStateType.Starting]):
+ runner_state: Literal[RunnerStateType.Starting]
+
+
+class DownloadingRunnerState(RunnerState[RunnerStateType.Downloading]):
+ runner_state: Literal[RunnerStateType.Downloading]
+ download_progress: BaseDownloadProgress[DownloadStatus]
+
+
+class RunningRunnerState(RunnerState[RunnerStateType.Running]):
+ runner_state: Literal[RunnerStateType.Running]
+
+
+class FailedRunnerState(RunnerState[RunnerStateType.Failed]):
+ runner_state: Literal[RunnerStateType.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)
+
+
+class ShardAssignments(BaseModel):
+ model_id: ModelId
+ runner_to_shard: Mapping[RunnerId, ShardMetadata[PartitionStrategy]]
+ node_to_runner: Mapping[NodeId, Sequence[RunnerId]]
+
+ @model_validator(mode="after")
+ def validate_runners_exist(self) -> "ShardAssignments":
+ for runners in self.node_to_runner.values():
+ for runner_id in runners:
+ if runner_id not in self.runner_to_shard:
+ raise ValueError(
+ f"Runner {runner_id} in node_to_runner does not exist in runner_to_shard"
+ )
+ return self
diff --git a/shared/types/worker/shards.py b/shared/types/worker/shards.py
new file mode 100644
index 00000000..5b33457d
--- /dev/null
+++ b/shared/types/worker/shards.py
@@ -0,0 +1,54 @@
+from enum import Enum
+from typing import Annotated, Generic, Literal, TypeVar
+
+from pydantic import BaseModel, DirectoryPath, Field, TypeAdapter
+
+from shared.types.common import NodeId
+from shared.types.models.common import ModelId
+
+
+class PartitionStrategy(str, Enum):
+ pipeline = "pipeline"
+
+
+PartitionStrategyT = TypeVar(name="PartitionStrategyT", bound=PartitionStrategy)
+
+
+class ShardMetadata(BaseModel, Generic[PartitionStrategyT]):
+ """
+ Defines a specific shard of the model that is ready to be run on a device.
+ Replaces previous `Shard` object.
+ """
+
+ device_rank: int
+ world_size: int
+ model_id: ModelId
+ model_path: DirectoryPath
+
+
+class PipelineShardMeta(ShardMetadata[PartitionStrategy.pipeline]):
+ """
+ Pipeline parallelism shard meta.
+ """
+
+ 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")]
+ShardMetaAdapter: TypeAdapter[ShardMetadata[PartitionStrategy]] = TypeAdapter(
+ _ShardMeta
+)
+
+
+class ShardPlacement(BaseModel, Generic[PartitionStrategyT]):
+ """
+ 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, ShardMetadata[PartitionStrategyT]]
diff --git a/shared/utils.py b/shared/utils.py
new file mode 100644
index 00000000..bf2be769
--- /dev/null
+++ b/shared/utils.py
@@ -0,0 +1,9 @@
+from typing import Any, Type, TypeVar
+
+T = TypeVar("T")
+
+
+def ensure_type(obj: Any, expected_type: Type[T]) -> T: # type: ignore
+ if not isinstance(obj, expected_type):
+ raise TypeError(f"Expected {expected_type}, got {type(obj)}") # type: ignore
+ return obj
diff --git a/uv.lock b/uv.lock
index 825473ce..d08efbb3 100644
--- a/uv.lock
+++ b/uv.lock
@@ -29,6 +29,19 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" },
]
+[[package]]
+name = "anyio"
+version = "4.9.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "idna", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+ { name = "sniffio", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/95/7d/4c1bd541d4dffa1b52bd83fb8527089e097a106fc90b467a7313b105f840/anyio-4.9.0.tar.gz", hash = "sha256:673c0c244e15788651a4ff38710fea9675823028a6f08a5eda409e0c9840a028", size = 190949, upload-time = "2025-03-17T00:02:54.77Z" }
+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.29.4"
@@ -41,6 +54,24 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/d5/dc/180fe721a2574fb3aad4051adcca196ac2d18adaf75122f5eeb47436cca2/basedpyright-1.29.4-py3-none-any.whl", hash = "sha256:e087513979972f83010639c6c1a1c13dd3b1d24ee45f8ecff747962cc2063d6f", size = 11476859, upload-time = "2025-06-11T22:25:52.01Z" },
]
+[[package]]
+name = "certifi"
+version = "2025.6.15"
+source = { registry = "https://pypi.org/simple" }
+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/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]]
+name = "distro"
+version = "1.9.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/fc/f8/98eea607f65de6527f8a2e8885fc8015d3e6f5775df186e443e0964a11c3/distro-1.9.0.tar.gz", hash = "sha256:2fa77c6fd8940f116ee1d6b94a2f90b13b5ea8d019b98bc8bafdcabcdd9bdbed", size = 60722, upload-time = "2023-12-24T09:54:32.31Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/12/b3/231ffd4ab1fc9d679809f356cebee130ac7daa00d6d6f3206dd4fd137e9e/distro-1.9.0-py3-none-any.whl", hash = "sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2", size = 20277, upload-time = "2023-12-24T09:54:30.421Z" },
+]
+
[[package]]
name = "exo"
version = "0.2.0"
@@ -105,6 +136,7 @@ name = "exo-shared"
version = "0.1.0"
source = { editable = "shared" }
dependencies = [
+ { 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'" },
@@ -118,6 +150,7 @@ dev = [
[package.metadata]
requires-dist = [
+ { name = "openai", specifier = ">=1.93.0" },
{ name = "pathlib", specifier = ">=1.0.1" },
{ name = "protobuf", specifier = ">=6.31.1" },
{ name = "pydantic", specifier = ">=2.11.7" },
@@ -138,6 +171,52 @@ dependencies = [
[package.metadata]
requires-dist = [{ name = "exo-shared", editable = "shared" }]
+[[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"
@@ -147,6 +226,38 @@ 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 = "jiter"
+version = "0.10.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/ee/9d/ae7ddb4b8ab3fb1b51faf4deb36cb48a4fbbd7cb36bad6a5fca4741306f7/jiter-0.10.0.tar.gz", hash = "sha256:07a7142c38aacc85194391108dc91b5b57093c978a9932bd86a36862759d9500", size = 162759, upload-time = "2025-05-18T19:04:59.73Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/2e/b0/279597e7a270e8d22623fea6c5d4eeac328e7d95c236ed51a2b884c54f70/jiter-0.10.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:e0588107ec8e11b6f5ef0e0d656fb2803ac6cf94a96b2b9fc675c0e3ab5e8644", size = 311617, upload-time = "2025-05-18T19:04:02.078Z" },
+ { url = "https://files.pythonhosted.org/packages/91/e3/0916334936f356d605f54cc164af4060e3e7094364add445a3bc79335d46/jiter-0.10.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:cafc4628b616dc32530c20ee53d71589816cf385dd9449633e910d596b1f5c8a", size = 318947, upload-time = "2025-05-18T19:04:03.347Z" },
+ { url = "https://files.pythonhosted.org/packages/6a/8e/fd94e8c02d0e94539b7d669a7ebbd2776e51f329bb2c84d4385e8063a2ad/jiter-0.10.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:520ef6d981172693786a49ff5b09eda72a42e539f14788124a07530f785c3ad6", size = 344618, upload-time = "2025-05-18T19:04:04.709Z" },
+ { url = "https://files.pythonhosted.org/packages/6f/b0/f9f0a2ec42c6e9c2e61c327824687f1e2415b767e1089c1d9135f43816bd/jiter-0.10.0-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:554dedfd05937f8fc45d17ebdf298fe7e0c77458232bcb73d9fbbf4c6455f5b3", size = 368829, upload-time = "2025-05-18T19:04:06.912Z" },
+ { url = "https://files.pythonhosted.org/packages/e8/57/5bbcd5331910595ad53b9fd0c610392ac68692176f05ae48d6ce5c852967/jiter-0.10.0-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5bc299da7789deacf95f64052d97f75c16d4fc8c4c214a22bf8d859a4288a1c2", size = 491034, upload-time = "2025-05-18T19:04:08.222Z" },
+ { url = "https://files.pythonhosted.org/packages/9b/be/c393df00e6e6e9e623a73551774449f2f23b6ec6a502a3297aeeece2c65a/jiter-0.10.0-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:5161e201172de298a8a1baad95eb85db4fb90e902353b1f6a41d64ea64644e25", size = 388529, upload-time = "2025-05-18T19:04:09.566Z" },
+ { url = "https://files.pythonhosted.org/packages/42/3e/df2235c54d365434c7f150b986a6e35f41ebdc2f95acea3036d99613025d/jiter-0.10.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2e2227db6ba93cb3e2bf67c87e594adde0609f146344e8207e8730364db27041", size = 350671, upload-time = "2025-05-18T19:04:10.98Z" },
+ { url = "https://files.pythonhosted.org/packages/c6/77/71b0b24cbcc28f55ab4dbfe029f9a5b73aeadaba677843fc6dc9ed2b1d0a/jiter-0.10.0-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:15acb267ea5e2c64515574b06a8bf393fbfee6a50eb1673614aa45f4613c0cca", size = 390864, upload-time = "2025-05-18T19:04:12.722Z" },
+ { url = "https://files.pythonhosted.org/packages/6a/d3/ef774b6969b9b6178e1d1e7a89a3bd37d241f3d3ec5f8deb37bbd203714a/jiter-0.10.0-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:901b92f2e2947dc6dfcb52fd624453862e16665ea909a08398dde19c0731b7f4", size = 522989, upload-time = "2025-05-18T19:04:14.261Z" },
+ { url = "https://files.pythonhosted.org/packages/0c/41/9becdb1d8dd5d854142f45a9d71949ed7e87a8e312b0bede2de849388cb9/jiter-0.10.0-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:d0cb9a125d5a3ec971a094a845eadde2db0de85b33c9f13eb94a0c63d463879e", size = 513495, upload-time = "2025-05-18T19:04:15.603Z" },
+ { url = "https://files.pythonhosted.org/packages/54/46/caa2c1342655f57d8f0f2519774c6d67132205909c65e9aa8255e1d7b4f4/jiter-0.10.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:28ed2a4c05a1f32ef0e1d24c2611330219fed727dae01789f4a335617634b1ca", size = 318225, upload-time = "2025-05-18T19:04:20.583Z" },
+ { url = "https://files.pythonhosted.org/packages/43/84/c7d44c75767e18946219ba2d703a5a32ab37b0bc21886a97bc6062e4da42/jiter-0.10.0-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:14a4c418b1ec86a195f1ca69da8b23e8926c752b685af665ce30777233dfe070", size = 350235, upload-time = "2025-05-18T19:04:22.363Z" },
+ { url = "https://files.pythonhosted.org/packages/1c/9b/1d646da42c3de6c2188fdaa15bce8ecb22b635904fc68be025e21249ba44/jiter-0.10.0-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:5e9251a5e83fab8d87799d3e1a46cb4b7f2919b895c6f4483629ed2446f66522", size = 310866, upload-time = "2025-05-18T19:04:24.891Z" },
+ { url = "https://files.pythonhosted.org/packages/ad/0e/26538b158e8a7c7987e94e7aeb2999e2e82b1f9d2e1f6e9874ddf71ebda0/jiter-0.10.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:023aa0204126fe5b87ccbcd75c8a0d0261b9abdbbf46d55e7ae9f8e22424eeb8", size = 318772, upload-time = "2025-05-18T19:04:26.161Z" },
+ { url = "https://files.pythonhosted.org/packages/7b/fb/d302893151caa1c2636d6574d213e4b34e31fd077af6050a9c5cbb42f6fb/jiter-0.10.0-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3c189c4f1779c05f75fc17c0c1267594ed918996a231593a21a5ca5438445216", size = 344534, upload-time = "2025-05-18T19:04:27.495Z" },
+ { url = "https://files.pythonhosted.org/packages/01/d8/5780b64a149d74e347c5128d82176eb1e3241b1391ac07935693466d6219/jiter-0.10.0-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:15720084d90d1098ca0229352607cd68256c76991f6b374af96f36920eae13c4", size = 369087, upload-time = "2025-05-18T19:04:28.896Z" },
+ { url = "https://files.pythonhosted.org/packages/e8/5b/f235a1437445160e777544f3ade57544daf96ba7e96c1a5b24a6f7ac7004/jiter-0.10.0-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:e4f2fb68e5f1cfee30e2b2a09549a00683e0fde4c6a2ab88c94072fc33cb7426", size = 490694, upload-time = "2025-05-18T19:04:30.183Z" },
+ { url = "https://files.pythonhosted.org/packages/85/a9/9c3d4617caa2ff89cf61b41e83820c27ebb3f7b5fae8a72901e8cd6ff9be/jiter-0.10.0-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:ce541693355fc6da424c08b7edf39a2895f58d6ea17d92cc2b168d20907dee12", size = 388992, upload-time = "2025-05-18T19:04:32.028Z" },
+ { url = "https://files.pythonhosted.org/packages/68/b1/344fd14049ba5c94526540af7eb661871f9c54d5f5601ff41a959b9a0bbd/jiter-0.10.0-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:31c50c40272e189d50006ad5c73883caabb73d4e9748a688b216e85a9a9ca3b9", size = 351723, upload-time = "2025-05-18T19:04:33.467Z" },
+ { url = "https://files.pythonhosted.org/packages/41/89/4c0e345041186f82a31aee7b9d4219a910df672b9fef26f129f0cda07a29/jiter-0.10.0-cp314-cp314-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:fa3402a2ff9815960e0372a47b75c76979d74402448509ccd49a275fa983ef8a", size = 392215, upload-time = "2025-05-18T19:04:34.827Z" },
+ { url = "https://files.pythonhosted.org/packages/55/58/ee607863e18d3f895feb802154a2177d7e823a7103f000df182e0f718b38/jiter-0.10.0-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:1956f934dca32d7bb647ea21d06d93ca40868b505c228556d3373cbd255ce853", size = 522762, upload-time = "2025-05-18T19:04:36.19Z" },
+ { url = "https://files.pythonhosted.org/packages/15/d0/9123fb41825490d16929e73c212de9a42913d68324a8ce3c8476cae7ac9d/jiter-0.10.0-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:fcedb049bdfc555e261d6f65a6abe1d5ad68825b7202ccb9692636c70fcced86", size = 513427, upload-time = "2025-05-18T19:04:37.544Z" },
+ { url = "https://files.pythonhosted.org/packages/03/0c/5fe86614ea050c3ecd728ab4035534387cd41e7c1855ef6c031f1ca93e3f/jiter-0.10.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:5ed975b83a2b8639356151cef5c0d597c68376fc4922b45d0eb384ac058cfa00", size = 318527, upload-time = "2025-05-18T19:04:40.612Z" },
+ { url = "https://files.pythonhosted.org/packages/b3/4a/4175a563579e884192ba6e81725fc0448b042024419be8d83aa8a80a3f44/jiter-0.10.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3aa96f2abba33dc77f79b4cf791840230375f9534e5fac927ccceb58c5e604a5", size = 354213, upload-time = "2025-05-18T19:04:41.894Z" },
+]
+
[[package]]
name = "markdown-it-py"
version = "3.0.0"
@@ -210,6 +321,25 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/4a/cd/e2b5083df581fc1d08eb93feb6f8fbd3d56b113cef9b59d8e0fb7d4dd4f3/nodejs_wheel_binaries-22.16.0-py2.py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:7f526ca6a132b0caf633566a2a78c6985fe92857e7bfdb37380f76205a10b808", size = 60763005, upload-time = "2025-05-22T07:27:41.39Z" },
]
+[[package]]
+name = "openai"
+version = "1.93.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "anyio", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+ { name = "distro", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+ { name = "httpx", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+ { name = "jiter", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+ { name = "pydantic", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+ { name = "sniffio", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+ { 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/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/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]]
name = "packaging"
version = "25.0"
@@ -347,6 +477,24 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/5a/c0/b0b508193b0e8a1654ec683ebab18d309861f8bd64e3a2f9648b80d392cb/ruff-0.11.13-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:51c3f95abd9331dc5b87c47ac7f376db5616041173826dfd556cfe3d4977f492", size = 11602992, upload-time = "2025-06-05T21:00:06.249Z" },
]
+[[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 = "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.20250516"
← 74d56e52 fix: Improve naming
·
back to Exo
·
New Runner! 21acd379 →