From b0bd95100501ef2f46091370510f12ae99c290d7 Mon Sep 17 00:00:00 2001 From: Arbion Halili <99731180+ToxicPine@users.noreply.github.com> Date: Wed, 9 Jul 2025 19:04:21 +0100 Subject: [PATCH] Merge Basic Interfaces Co-authored-by: Alex Cheema Co-authored-by: Seth Howes Co-authored-by: Matt Beton Co-authored-by: Andrei Cravtov --- .../{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(-) rename .github/actions/conditional-commit/{conditional-commit.yml => action.yml} (100%) create mode 100644 .github/actions/lint-check/action.yml create mode 100644 .github/actions/verify-clean/action.yml create mode 100644 master/api.py create mode 100644 shared/openai.py create mode 100644 shared/protobufs/types/mlx/nn/__init__.pyi create mode 100644 shared/types/api.py create mode 100644 shared/types/common.py delete mode 100644 shared/types/event_sourcing.py create mode 100644 shared/types/events/chunks.py create mode 100644 shared/types/events/common.py create mode 100644 shared/types/events/events.py create mode 100644 shared/types/graphs/common.py create mode 100644 shared/types/graphs/resource_graph.py create mode 100644 shared/types/models/common.py create mode 100644 shared/types/models/metadata.py create mode 100644 shared/types/models/model.py create mode 100644 shared/types/models/sources.py create mode 100644 shared/types/networking/control_plane.py create mode 100644 shared/types/networking/data_plane.py create mode 100644 shared/types/networking/services.py create mode 100644 shared/types/networking/topology.py create mode 100644 shared/types/profiling/common.py create mode 100644 shared/types/states/master.py create mode 100644 shared/types/states/shared.py create mode 100644 shared/types/states/worker.py create mode 100644 shared/types/tasks/common.py create mode 100644 shared/types/worker/commands_runner.py create mode 100644 shared/types/worker/common.py create mode 100644 shared/types/worker/downloads.py create mode 100644 shared/types/worker/instances.py create mode 100644 shared/types/worker/mlx.py create mode 100644 shared/types/worker/resource_monitor.py create mode 100644 shared/types/worker/runners.py create mode 100644 shared/types/worker/shards.py create mode 100644 shared/utils.py 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 + - uses: ./.github/actions/verify-clean with: - message: "chore(proto) regenerate protobufs" + step: regenerate-protobufs - - uses: ./.github/actions/format - - - name: Commit formatted code - uses: ./.github/actions/conditional-commit - with: - message: "chore(format): format code" - - - uses: ./.github/actions/lint - - - 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"