From bb7f1ae99421f3e9d075c0a9ffbb6e5ce9faf8dd Mon Sep 17 00:00:00 2001 From: Alex Cheema <41707476+AlexCheema@users.noreply.github.com> Date: Fri, 18 Jul 2025 10:08:56 +0100 Subject: [PATCH] New worker Co-authored-by: Matt Beton --- .github/actions/typecheck/action.yml | 4 +- engines/mlx/auto_parallel.py | 1 - engines/mlx/utils_mlx.py | 8 +- shared/types/api.py | 3 +- shared/types/events/common.py | 12 +- shared/types/events/events.py | 87 +++--- shared/types/states/master.py | 6 +- shared/types/states/shared.py | 4 +- shared/types/states/worker.py | 4 + shared/types/tasks/common.py | 93 +++++-- shared/types/worker/commands_runner.py | 9 +- shared/types/worker/downloads.py | 12 +- shared/types/worker/instances.py | 2 + shared/types/worker/mlx.py | 6 +- shared/types/worker/ops.py | 69 +++++ shared/types/worker/runners.py | 78 +++--- shared/types/worker/shards.py | 30 ++- test_shard_serialization.py | 1 + worker/main.py | 358 +++++++++++++++++++++++++ worker/runner/communication.py | 4 +- worker/runner/runner.py | 44 ++- worker/runner/runner_supervisor.py | 29 +- worker/test_worker_state.py | 48 ++++ worker/tests/conftest.py | 155 ++++++++--- worker/tests/test_serdes.py | 8 +- worker/tests/test_supervisor.py | 55 ++-- worker/tests/test_worker_handlers.py | 211 +++++++++++++++ worker/tests/test_worker_plan.py | 263 ++++++++++++++++++ 28 files changed, 1368 insertions(+), 236 deletions(-) create mode 100644 shared/types/worker/ops.py create mode 100644 test_shard_serialization.py create mode 100644 worker/test_worker_state.py create mode 100644 worker/tests/test_worker_handlers.py create mode 100644 worker/tests/test_worker_plan.py diff --git a/.github/actions/typecheck/action.yml b/.github/actions/typecheck/action.yml index 8ae7ffa2..ba61737f 100644 --- a/.github/actions/typecheck/action.yml +++ b/.github/actions/typecheck/action.yml @@ -6,5 +6,7 @@ runs: using: "composite" steps: - name: Run type checker - run: nix develop -c just check + run: | + nix develop -c just sync + nix develop -c just check shell: bash diff --git a/engines/mlx/auto_parallel.py b/engines/mlx/auto_parallel.py index 3b8531bb..a75d356e 100644 --- a/engines/mlx/auto_parallel.py +++ b/engines/mlx/auto_parallel.py @@ -79,7 +79,6 @@ def auto_parallel( Returns: The parallelized model """ - inner_model_instance: nn.Module = inner_model(model) # Handle both model.layers and model.h cases diff --git a/engines/mlx/utils_mlx.py b/engines/mlx/utils_mlx.py index 5de40e63..bae55498 100644 --- a/engines/mlx/utils_mlx.py +++ b/engines/mlx/utils_mlx.py @@ -14,8 +14,8 @@ from mlx_lm.tokenizer_utils import TokenizerWrapper, load_tokenizer from mlx_lm.utils import load_model from pydantic import RootModel -from shared.mlx.auto_parallel import auto_parallel -from shared.types.tasks.common import ChatCompletionParams +from engines.mlx.auto_parallel import auto_parallel +from shared.types.tasks.common import CompletionCreateParams from shared.types.worker.mlx import Host from shared.types.worker.shards import ShardMeta from worker.runner.communication import runner_print @@ -96,12 +96,12 @@ def shard_and_load(model_shard_meta: ShardMeta) -> tuple[nn.Module, TokenizerWra async def apply_chat_template( mlx_executor: concurrent.futures.ThreadPoolExecutor, tokenizer: TokenizerWrapper, - chat_task: ChatCompletionParams, + chat_task_data: CompletionCreateParams, ) -> str: loop: AbstractEventLoop = asyncio.get_running_loop() # Now we can properly access the messages - messages = chat_task.messages + messages = chat_task_data.messages messages_dicts = [msg.model_dump() for msg in messages] # Filter out None values, keeping only 'role' and 'content' keys diff --git a/shared/types/api.py b/shared/types/api.py index f1bdefbf..5bf878ef 100644 --- a/shared/types/api.py +++ b/shared/types/api.py @@ -1,9 +1,8 @@ from typing import Literal -from openai.types.chat.completion_create_params import CompletionCreateParams from pydantic import BaseModel -from shared.types.tasks.common import TaskId +from shared.types.tasks.common import CompletionCreateParams, TaskId class ChatTask(BaseModel): diff --git a/shared/types/events/common.py b/shared/types/events/common.py index b4c3ae40..0c825c21 100644 --- a/shared/types/events/common.py +++ b/shared/types/events/common.py @@ -133,10 +133,16 @@ EventCategories = FrozenSet[EventCategory] assert_literal_union_covers_enum(EventCategory, EventCategoryEnum) -class BaseEvent[SetMembersT: EventCategories | EventCategory](BaseModel): - event_type: EventTypes +EventTypeT = EventTypes # Type Alias placeholder; generic parameter will override + + +class BaseEvent[ + SetMembersT: EventCategories | EventCategory, + EventTypeLitT: EventTypes = EventTypes, +](BaseModel): + event_type: EventTypeLitT event_category: SetMembersT - event_id: EventId + event_id: EventId = EventId() def check_event_was_sent_by_correct_node(self, origin_id: NodeId) -> bool: ... diff --git a/shared/types/events/events.py b/shared/types/events/events.py index 1c52f59f..f7a609b4 100644 --- a/shared/types/events/events.py +++ b/shared/types/events/events.py @@ -9,7 +9,6 @@ from shared.types.events.common import ( ControlPlaneEventTypes, DataPlaneEventTypes, EventCategoryEnum, - EventTypes, InstanceEventTypes, NodePerformanceEventTypes, RunnerStatusEventTypes, @@ -28,8 +27,8 @@ from shared.types.networking.data_plane import ( ) from shared.types.profiling.common import NodePerformanceProfile from shared.types.tasks.common import ( + BaseTaskData, TaskId, - TaskParams, TaskState, TaskStatusOtherType, TaskStatusType, @@ -37,7 +36,7 @@ from shared.types.tasks.common import ( ) from shared.types.worker.common import InstanceId, NodeStatus from shared.types.worker.instances import InstanceParams, TypeOfInstance -from shared.types.worker.runners import RunnerId, RunnerStatus, RunnerStatusType +from shared.types.worker.runners import RunnerId, RunnerStatus TaskEvent = BaseEvent[EventCategoryEnum.MutatesTaskState] InstanceEvent = BaseEvent[EventCategoryEnum.MutatesInstanceState] @@ -46,132 +45,132 @@ DataPlaneEvent = BaseEvent[EventCategoryEnum.MutatesDataPlaneState] NodePerformanceEvent = BaseEvent[EventCategoryEnum.MutatesNodePerformanceState] -class TaskCreated(BaseEvent[EventCategoryEnum.MutatesTaskState]): - event_type: EventTypes = TaskEventTypes.TaskCreated +class TaskCreated(BaseEvent[EventCategoryEnum.MutatesTaskState, Literal[TaskEventTypes.TaskCreated]]): + event_type: Literal[TaskEventTypes.TaskCreated] = TaskEventTypes.TaskCreated event_category: Literal[EventCategoryEnum.MutatesTaskState] = EventCategoryEnum.MutatesTaskState task_id: TaskId - task_params: TaskParams[TaskType] + task_data: BaseTaskData[TaskType] task_state: TaskState[Literal[TaskStatusOtherType.Pending], TaskType] on_instance: InstanceId # Covers Cancellation Of Task, Non-Cancelled Tasks Perist -class TaskDeleted(BaseEvent[EventCategoryEnum.MutatesTaskState]): - event_type: EventTypes = TaskEventTypes.TaskDeleted +class TaskDeleted(BaseEvent[EventCategoryEnum.MutatesTaskState, Literal[TaskEventTypes.TaskDeleted]]): + event_type: Literal[TaskEventTypes.TaskDeleted] = TaskEventTypes.TaskDeleted event_category: Literal[EventCategoryEnum.MutatesTaskState] = EventCategoryEnum.MutatesTaskState task_id: TaskId -class TaskStateUpdated(BaseEvent[EventCategoryEnum.MutatesTaskState]): - event_type: EventTypes = TaskEventTypes.TaskStateUpdated +class TaskStateUpdated(BaseEvent[EventCategoryEnum.MutatesTaskState, Literal[TaskEventTypes.TaskStateUpdated]]): + event_type: Literal[TaskEventTypes.TaskStateUpdated] = TaskEventTypes.TaskStateUpdated event_category: Literal[EventCategoryEnum.MutatesTaskState] = EventCategoryEnum.MutatesTaskState task_state: TaskState[TaskStatusType, TaskType] -class InstanceCreated(BaseEvent[EventCategoryEnum.MutatesInstanceState]): - event_type: EventTypes = InstanceEventTypes.InstanceCreated +class InstanceCreated(BaseEvent[EventCategoryEnum.MutatesInstanceState, Literal[InstanceEventTypes.InstanceCreated]]): + event_type: Literal[InstanceEventTypes.InstanceCreated] = InstanceEventTypes.InstanceCreated event_category: Literal[EventCategoryEnum.MutatesInstanceState] = EventCategoryEnum.MutatesInstanceState instance_id: InstanceId instance_params: InstanceParams instance_type: TypeOfInstance -class InstanceActivated(BaseEvent[EventCategoryEnum.MutatesInstanceState]): - event_type: EventTypes = InstanceEventTypes.InstanceActivated +class InstanceActivated(BaseEvent[EventCategoryEnum.MutatesInstanceState, Literal[InstanceEventTypes.InstanceActivated]]): + event_type: Literal[InstanceEventTypes.InstanceActivated] = InstanceEventTypes.InstanceActivated event_category: Literal[EventCategoryEnum.MutatesInstanceState] = EventCategoryEnum.MutatesInstanceState instance_id: InstanceId -class InstanceDeactivated(BaseEvent[EventCategoryEnum.MutatesInstanceState]): - event_type: EventTypes = InstanceEventTypes.InstanceDeactivated +class InstanceDeactivated(BaseEvent[EventCategoryEnum.MutatesInstanceState, Literal[InstanceEventTypes.InstanceDeactivated]]): + event_type: Literal[InstanceEventTypes.InstanceDeactivated] = InstanceEventTypes.InstanceDeactivated event_category: Literal[EventCategoryEnum.MutatesInstanceState] = EventCategoryEnum.MutatesInstanceState instance_id: InstanceId -class InstanceDeleted(BaseEvent[EventCategoryEnum.MutatesInstanceState]): - event_type: EventTypes = InstanceEventTypes.InstanceDeleted +class InstanceDeleted(BaseEvent[EventCategoryEnum.MutatesInstanceState, Literal[InstanceEventTypes.InstanceDeleted]]): + event_type: Literal[InstanceEventTypes.InstanceDeleted] = InstanceEventTypes.InstanceDeleted event_category: Literal[EventCategoryEnum.MutatesInstanceState] = EventCategoryEnum.MutatesInstanceState instance_id: InstanceId transition: Tuple[InstanceId, InstanceId] -class InstanceReplacedAtomically(BaseEvent[EventCategoryEnum.MutatesInstanceState]): - event_type: EventTypes = InstanceEventTypes.InstanceReplacedAtomically +class InstanceReplacedAtomically(BaseEvent[EventCategoryEnum.MutatesInstanceState, Literal[InstanceEventTypes.InstanceReplacedAtomically]]): + event_type: Literal[InstanceEventTypes.InstanceReplacedAtomically] = InstanceEventTypes.InstanceReplacedAtomically event_category: Literal[EventCategoryEnum.MutatesInstanceState] = EventCategoryEnum.MutatesInstanceState instance_to_replace: InstanceId new_instance_id: InstanceId -class RunnerStatusUpdated(BaseEvent[EventCategoryEnum.MutatesRunnerStatus]): - event_type: EventTypes = RunnerStatusEventTypes.RunnerStatusUpdated +class RunnerStatusUpdated(BaseEvent[EventCategoryEnum.MutatesRunnerStatus, Literal[RunnerStatusEventTypes.RunnerStatusUpdated]]): + event_type: Literal[RunnerStatusEventTypes.RunnerStatusUpdated] = RunnerStatusEventTypes.RunnerStatusUpdated event_category: Literal[EventCategoryEnum.MutatesRunnerStatus] = EventCategoryEnum.MutatesRunnerStatus - instance_id: InstanceId - state_update: Tuple[RunnerId, RunnerStatus[RunnerStatusType]] + runner_id: RunnerId + runner_status: RunnerStatus -class MLXInferenceSagaPrepare(BaseEvent[EventCategoryEnum.MutatesTaskSagaState]): - event_type: EventTypes = TaskSagaEventTypes.MLXInferenceSagaPrepare +class MLXInferenceSagaPrepare(BaseEvent[EventCategoryEnum.MutatesTaskSagaState, Literal[TaskSagaEventTypes.MLXInferenceSagaPrepare]]): + event_type: Literal[TaskSagaEventTypes.MLXInferenceSagaPrepare] = TaskSagaEventTypes.MLXInferenceSagaPrepare event_category: Literal[EventCategoryEnum.MutatesTaskSagaState] = EventCategoryEnum.MutatesTaskSagaState task_id: TaskId instance_id: InstanceId -class MLXInferenceSagaStartPrepare(BaseEvent[EventCategoryEnum.MutatesTaskSagaState]): - event_type: EventTypes = TaskSagaEventTypes.MLXInferenceSagaStartPrepare +class MLXInferenceSagaStartPrepare(BaseEvent[EventCategoryEnum.MutatesTaskSagaState, Literal[TaskSagaEventTypes.MLXInferenceSagaStartPrepare]]): + event_type: Literal[TaskSagaEventTypes.MLXInferenceSagaStartPrepare] = TaskSagaEventTypes.MLXInferenceSagaStartPrepare event_category: Literal[EventCategoryEnum.MutatesTaskSagaState] = EventCategoryEnum.MutatesTaskSagaState task_id: TaskId instance_id: InstanceId -class NodePerformanceMeasured(BaseEvent[EventCategoryEnum.MutatesNodePerformanceState]): - event_type: EventTypes = NodePerformanceEventTypes.NodePerformanceMeasured +class NodePerformanceMeasured(BaseEvent[EventCategoryEnum.MutatesNodePerformanceState, Literal[NodePerformanceEventTypes.NodePerformanceMeasured]]): + event_type: Literal[NodePerformanceEventTypes.NodePerformanceMeasured] = NodePerformanceEventTypes.NodePerformanceMeasured event_category: Literal[EventCategoryEnum.MutatesNodePerformanceState] = EventCategoryEnum.MutatesNodePerformanceState node_id: NodeId node_profile: NodePerformanceProfile -class WorkerConnected(BaseEvent[EventCategoryEnum.MutatesControlPlaneState]): - event_type: EventTypes = ControlPlaneEventTypes.WorkerConnected +class WorkerConnected(BaseEvent[EventCategoryEnum.MutatesControlPlaneState, Literal[ControlPlaneEventTypes.WorkerConnected]]): + event_type: Literal[ControlPlaneEventTypes.WorkerConnected] = ControlPlaneEventTypes.WorkerConnected event_category: Literal[EventCategoryEnum.MutatesControlPlaneState] = EventCategoryEnum.MutatesControlPlaneState edge: DataPlaneEdge -class WorkerStatusUpdated(BaseEvent[EventCategoryEnum.MutatesControlPlaneState]): - event_type: EventTypes = ControlPlaneEventTypes.WorkerStatusUpdated +class WorkerStatusUpdated(BaseEvent[EventCategoryEnum.MutatesControlPlaneState, Literal[ControlPlaneEventTypes.WorkerStatusUpdated]]): + event_type: Literal[ControlPlaneEventTypes.WorkerStatusUpdated] = ControlPlaneEventTypes.WorkerStatusUpdated event_category: Literal[EventCategoryEnum.MutatesControlPlaneState] = EventCategoryEnum.MutatesControlPlaneState node_id: NodeId node_state: NodeStatus -class WorkerDisconnected(BaseEvent[EventCategoryEnum.MutatesControlPlaneState]): - event_type: EventTypes = ControlPlaneEventTypes.WorkerConnected +class WorkerDisconnected(BaseEvent[EventCategoryEnum.MutatesControlPlaneState, Literal[ControlPlaneEventTypes.WorkerDisconnected]]): + event_type: Literal[ControlPlaneEventTypes.WorkerDisconnected] = ControlPlaneEventTypes.WorkerDisconnected event_category: Literal[EventCategoryEnum.MutatesControlPlaneState] = EventCategoryEnum.MutatesControlPlaneState vertex_id: ControlPlaneEdgeId -class ChunkGenerated(BaseEvent[EventCategoryEnum.MutatesTaskState]): - event_type: EventTypes = StreamingEventTypes.ChunkGenerated +class ChunkGenerated(BaseEvent[EventCategoryEnum.MutatesTaskState, Literal[StreamingEventTypes.ChunkGenerated]]): + event_type: Literal[StreamingEventTypes.ChunkGenerated] = StreamingEventTypes.ChunkGenerated event_category: Literal[EventCategoryEnum.MutatesTaskState] = EventCategoryEnum.MutatesTaskState task_id: TaskId chunk: GenerationChunk -class DataPlaneEdgeCreated(BaseEvent[EventCategoryEnum.MutatesDataPlaneState]): - event_type: EventTypes = DataPlaneEventTypes.DataPlaneEdgeCreated +class DataPlaneEdgeCreated(BaseEvent[EventCategoryEnum.MutatesDataPlaneState, Literal[DataPlaneEventTypes.DataPlaneEdgeCreated]]): + event_type: Literal[DataPlaneEventTypes.DataPlaneEdgeCreated] = DataPlaneEventTypes.DataPlaneEdgeCreated event_category: Literal[EventCategoryEnum.MutatesDataPlaneState] = EventCategoryEnum.MutatesDataPlaneState vertex: ControlPlaneEdgeType -class DataPlaneEdgeReplacedAtomically(BaseEvent[EventCategoryEnum.MutatesDataPlaneState]): - event_type: EventTypes = DataPlaneEventTypes.DataPlaneEdgeReplacedAtomically +class DataPlaneEdgeReplacedAtomically(BaseEvent[EventCategoryEnum.MutatesDataPlaneState, Literal[DataPlaneEventTypes.DataPlaneEdgeReplacedAtomically]]): + event_type: Literal[DataPlaneEventTypes.DataPlaneEdgeReplacedAtomically] = DataPlaneEventTypes.DataPlaneEdgeReplacedAtomically event_category: Literal[EventCategoryEnum.MutatesDataPlaneState] = EventCategoryEnum.MutatesDataPlaneState edge_id: DataPlaneEdgeId edge_profile: DataPlaneEdgeProfile -class DataPlaneEdgeDeleted(BaseEvent[EventCategoryEnum.MutatesDataPlaneState]): - event_type: EventTypes = DataPlaneEventTypes.DataPlaneEdgeDeleted +class DataPlaneEdgeDeleted(BaseEvent[EventCategoryEnum.MutatesDataPlaneState, Literal[DataPlaneEventTypes.DataPlaneEdgeDeleted]]): + event_type: Literal[DataPlaneEventTypes.DataPlaneEdgeDeleted] = DataPlaneEventTypes.DataPlaneEdgeDeleted event_category: Literal[EventCategoryEnum.MutatesDataPlaneState] = EventCategoryEnum.MutatesDataPlaneState edge_id: DataPlaneEdgeId diff --git a/shared/types/states/master.py b/shared/types/states/master.py index 8a078d09..46a7348e 100644 --- a/shared/types/states/master.py +++ b/shared/types/states/master.py @@ -26,7 +26,7 @@ from shared.types.networking.topology import ( ) from shared.types.profiling.common import NodePerformanceProfile from shared.types.states.shared import SharedState -from shared.types.tasks.common import TaskParams, TaskType +from shared.types.tasks.common import BaseTaskData, TaskType from shared.types.worker.common import NodeStatus from shared.types.worker.instances import InstanceId, InstanceParams @@ -78,8 +78,8 @@ class ControlPlaneNetworkState(State[EventCategoryEnum.MutatesControlPlaneState] class MasterState(SharedState): data_plane_network_state: DataPlaneNetworkState = DataPlaneNetworkState() control_plane_network_state: ControlPlaneNetworkState = ControlPlaneNetworkState() - job_inbox: Queue[TaskParams[TaskType]] = Queue() - job_outbox: Queue[TaskParams[TaskType]] = Queue() + job_inbox: Queue[BaseTaskData[TaskType]] = Queue() + job_outbox: Queue[BaseTaskData[TaskType]] = Queue() cache_policy: CachePolicy[CachePolicyType] = CachePolicy[CachePolicyType]( policy_type=CachePolicyType.KeepAll ) diff --git a/shared/types/states/shared.py b/shared/types/states/shared.py index 388e1cbe..58b4331a 100644 --- a/shared/types/states/shared.py +++ b/shared/types/states/shared.py @@ -14,7 +14,7 @@ from shared.types.tasks.common import ( ) from shared.types.worker.common import InstanceId from shared.types.worker.instances import BaseInstance -from shared.types.worker.runners import RunnerId, RunnerStatus, RunnerStatusType +from shared.types.worker.runners import RunnerId, RunnerStatus class Instances(State[EventCategoryEnum.MutatesInstanceState]): @@ -42,7 +42,7 @@ class Runners(State[EventCategoryEnum.MutatesRunnerStatus]): event_category: Literal[EventCategoryEnum.MutatesRunnerStatus] = ( EventCategoryEnum.MutatesRunnerStatus ) - runner_statuses: Mapping[RunnerId, RunnerStatus[RunnerStatusType]] = {} + runner_statuses: Mapping[RunnerId, RunnerStatus] = {} class SharedState(BaseModel): diff --git a/shared/types/states/worker.py b/shared/types/states/worker.py index dfddc265..6fdef1a8 100644 --- a/shared/types/states/worker.py +++ b/shared/types/states/worker.py @@ -1,4 +1,5 @@ from collections.abc import Mapping +from typing import Literal from shared.types.common import NodeId from shared.types.events.common import ( @@ -10,6 +11,9 @@ from shared.types.worker.common import NodeStatus class NodeStatusState(State[EventCategoryEnum.MutatesControlPlaneState]): + event_category: Literal[EventCategoryEnum.MutatesControlPlaneState] = ( + EventCategoryEnum.MutatesControlPlaneState + ) node_status: Mapping[NodeId, NodeStatus] diff --git a/shared/types/tasks/common.py b/shared/types/tasks/common.py index 2b422d6e..42468d4f 100644 --- a/shared/types/tasks/common.py +++ b/shared/types/tasks/common.py @@ -1,7 +1,15 @@ from enum import Enum -from typing import Annotated, Generic, Literal, TypeVar, Union, final +from typing import ( # noqa: E402 + Annotated, + Any, + Generic, + Literal, + TypeAlias, + TypeVar, + Union, + final, +) -import openai.types.chat as openai from pydantic import BaseModel, Field, TypeAdapter from shared.types.common import NewUUID @@ -12,34 +20,14 @@ class TaskId(NewUUID): pass +## TASK TYPES @final class TaskType(str, Enum): - ChatCompletionNonStreaming = "ChatCompletionNonStreaming" - ChatCompletionStreaming = "ChatCompletionStreaming" - + ChatCompletion = "ChatCompletion" TaskTypeT = TypeVar("TaskTypeT", bound=TaskType, covariant=True) - -class TaskParams(BaseModel, Generic[TaskTypeT]): ... - - -@final -class ChatCompletionNonStreamingTask(TaskParams[TaskType.ChatCompletionNonStreaming]): - task_type: Literal[TaskType.ChatCompletionNonStreaming] = ( - TaskType.ChatCompletionNonStreaming - ) - task_data: openai.completion_create_params.CompletionCreateParamsNonStreaming - - -@final -class ChatCompletionStreamingTask(TaskParams[TaskType.ChatCompletionStreaming]): - task_type: Literal[TaskType.ChatCompletionStreaming] = ( - TaskType.ChatCompletionStreaming - ) - task_data: openai.completion_create_params.CompletionCreateParamsStreaming - - +## TASK STATUSES @final class TaskStatusFailedType(str, Enum): Failed = "Failed" @@ -57,8 +45,56 @@ class TaskStatusOtherType(str, Enum): TaskStatusType = TaskStatusCompleteType | TaskStatusFailedType | TaskStatusOtherType +TaskStatusTypeT = TypeVar("TaskStatusTypeT", bound=TaskStatusType)#, covariant=True) +## Peripherals +class ChatCompletionMessage(BaseModel): + role: Literal["system", "user", "assistant", "developer", "tool", "function"] + content: str | None = None + name: str | None = None + tool_calls: list[dict[str, Any]] | None = None + tool_call_id: str | None = None + function_call: dict[str, Any] | None = None + +class CompletionCreateParams(BaseModel): + model: str + messages: list[ChatCompletionMessage] + frequency_penalty: float | None = None + logit_bias: dict[str, int] | None = None + logprobs: bool | None = None + top_logprobs: int | None = None + max_tokens: int | None = None + n: int | None = None + presence_penalty: float | None = None + response_format: dict[str, Any] | None = None + seed: int | None = None + stop: str | list[str] | None = None + stream: bool = False + temperature: float | None = None + top_p: float | None = None + tools: list[dict[str, Any]] | None = None + tool_choice: str | dict[str, Any] | None = None + parallel_tool_calls: bool | None = None + user: str | None = None + + +## Task Data is stored in task, one-to-one with task type + +class BaseTaskData(BaseModel, Generic[TaskTypeT]): ... + +@final +class ChatCompletionTaskData(BaseTaskData[TaskType.ChatCompletion]): + task_type: Literal[TaskType.ChatCompletion] = ( + TaskType.ChatCompletion + ) + task_params: CompletionCreateParams + +TaskData: TypeAlias = ChatCompletionTaskData + + +## TASKS + class TaskArtifact[TaskTypeT: TaskType, TaskStatusTypeT: TaskStatusType](BaseModel): ... @@ -82,15 +118,14 @@ class TaskState[TaskStatusTypeT: TaskStatusType, TaskTypeT: TaskType](BaseModel) class BaseTask[TaskTypeT: TaskType, TaskStatusTypeT: TaskStatusType](BaseModel): task_type: TaskTypeT - task_params: TaskParams[TaskTypeT] + task_data: TaskData # Really this should be BaseTaskData[TaskTypeT], but this causes a bunch of errors that I don't know how to fix yet. task_state: TaskState[TaskStatusTypeT, TaskTypeT] on_instance: InstanceId BaseTaskAnnotated = Annotated[ Union[ - BaseTask[Literal[TaskType.ChatCompletionNonStreaming], TaskStatusType], - BaseTask[Literal[TaskType.ChatCompletionStreaming], TaskStatusType], + BaseTask[Literal[TaskType.ChatCompletion], TaskStatusType], ], Field(discriminator="task_type"), ] @@ -109,4 +144,4 @@ class TaskSagaEntry(BaseModel): class Task[TaskTypeT: TaskType, TaskStatusTypeT: TaskStatusType]( BaseTask[TaskTypeT, TaskStatusTypeT] ): - task_id: TaskId + task_id: TaskId \ No newline at end of file diff --git a/shared/types/worker/commands_runner.py b/shared/types/worker/commands_runner.py index 7f636588..83283135 100644 --- a/shared/types/worker/commands_runner.py +++ b/shared/types/worker/commands_runner.py @@ -4,9 +4,9 @@ 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.tasks.common import ChatCompletionTaskData from shared.types.worker.mlx import Host -from shared.types.worker.shards import PartitionStrategy, ShardMetadata +from shared.types.worker.shards import ShardMetadata ## Messages passed TO the runner @@ -26,15 +26,16 @@ class BaseRunnerMessage(BaseModel, Generic[MT]): class SetupMessage(BaseRunnerMessage[MessageType.Setup]): type: Literal[MessageType.Setup] = Field(default=MessageType.Setup, frozen=True) - model_shard_meta: ShardMetadata[PartitionStrategy] + model_shard_meta: ShardMetadata hosts: list[Host] +# TODO: We probably want a general task message that can take any task type. Can be fixed later. class ChatTaskMessage(BaseRunnerMessage[MessageType.ChatTask]): type: Literal[MessageType.ChatTask] = Field( default=MessageType.ChatTask, frozen=True ) - task: ChatTask + task_data: ChatCompletionTaskData class ExitMessage(BaseRunnerMessage[MessageType.Exit]): diff --git a/shared/types/worker/downloads.py b/shared/types/worker/downloads.py index acc53650..649eb48b 100644 --- a/shared/types/worker/downloads.py +++ b/shared/types/worker/downloads.py @@ -13,7 +13,7 @@ from pydantic import BaseModel, Field, PositiveInt from shared.types.common import NodeId from shared.types.models.common import ModelId from shared.types.models.sources import ModelSource -from shared.types.worker.shards import PartitionStrategy, ShardMetadata +from shared.types.worker.shards import ShardMetadata class DownloadProgressData(BaseModel): @@ -34,21 +34,21 @@ class BaseDownloadProgress[DownloadStatusT: DownloadStatus](BaseModel): class DownloadPending(BaseDownloadProgress[DownloadStatus.Pending]): - download_status: Literal[DownloadStatus.Pending] = Field(DownloadStatus.Pending) + download_status: Literal[DownloadStatus.Pending] = Field(default=DownloadStatus.Pending) class DownloadCompleted(BaseDownloadProgress[DownloadStatus.Completed]): - download_status: Literal[DownloadStatus.Completed] = Field(DownloadStatus.Completed) + download_status: Literal[DownloadStatus.Completed] = Field(default=DownloadStatus.Completed) class DownloadFailed(BaseDownloadProgress[DownloadStatus.Failed]): - download_status: Literal[DownloadStatus.Failed] = Field(DownloadStatus.Failed) + download_status: Literal[DownloadStatus.Failed] = Field(default=DownloadStatus.Failed) error_message: str class DownloadOngoing(BaseDownloadProgress[DownloadStatus.Downloading]): download_status: Literal[DownloadStatus.Downloading] = Field( - DownloadStatus.Downloading + default=DownloadStatus.Downloading ) download_progress: DownloadProgressData @@ -75,6 +75,6 @@ DownloadEffectHandler = Callable[ def download_shard( model_id: ModelId, model_source: ModelSource, - shard_metadata: ShardMetadata[PartitionStrategy], + shard_metadata: ShardMetadata, effect_handlers: Sequence[DownloadEffectHandler], ) -> None: ... diff --git a/shared/types/worker/instances.py b/shared/types/worker/instances.py index c3dd7bb8..50047adc 100644 --- a/shared/types/worker/instances.py +++ b/shared/types/worker/instances.py @@ -3,6 +3,7 @@ from enum import Enum from pydantic import BaseModel from shared.types.worker.common import InstanceId +from shared.types.worker.mlx import Host from shared.types.worker.runners import ( ShardAssignments, ) @@ -15,6 +16,7 @@ class TypeOfInstance(str, Enum): class InstanceParams(BaseModel): shard_assignments: ShardAssignments + hosts: list[Host] class BaseInstance(BaseModel): diff --git a/shared/types/worker/mlx.py b/shared/types/worker/mlx.py index 496ef369..9e8267bc 100644 --- a/shared/types/worker/mlx.py +++ b/shared/types/worker/mlx.py @@ -6,8 +6,12 @@ class Host(BaseModel): host: str port: int + def __str__(self) -> str: + return f"{self.host}:{self.port}" + @field_validator("port") - def check_port(self, v: int) -> int: + @classmethod + def check_port(cls, 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/ops.py b/shared/types/worker/ops.py new file mode 100644 index 00000000..5e0a9753 --- /dev/null +++ b/shared/types/worker/ops.py @@ -0,0 +1,69 @@ +from enum import Enum +from typing import Annotated, Generic, Literal, TypeVar, Union + +from pydantic import BaseModel, Field + +from shared.types.events.events import InstanceId +from shared.types.tasks.common import Task, TaskStatusType, TaskType +from shared.types.worker.common import RunnerId +from shared.types.worker.mlx import Host +from shared.types.worker.shards import ShardMetadata + + +class RunnerOpType(str, Enum): + ASSIGN_RUNNER = "assign_runner" + UNASSIGN_RUNNER = "unassign_runner" + RUNNER_UP = "runner_up" + RUNNER_DOWN = "runner_down" + DOWNLOAD = "download" + CHAT_COMPLETION = "chat_completion" + +RunnerOpT = TypeVar("RunnerOpT", bound=RunnerOpType) + +class BaseRunnerOp(BaseModel, Generic[RunnerOpT]): + op_type: RunnerOpT + +class AssignRunnerOp(BaseRunnerOp[Literal[RunnerOpType.ASSIGN_RUNNER]]): + op_type: Literal[RunnerOpType.ASSIGN_RUNNER] = Field(default=RunnerOpType.ASSIGN_RUNNER, frozen=True) + instance_id: InstanceId + runner_id: RunnerId + shard_metadata: ShardMetadata + hosts: list[Host] + +class UnassignRunnerOp(BaseRunnerOp[Literal[RunnerOpType.UNASSIGN_RUNNER]]): + op_type: Literal[RunnerOpType.UNASSIGN_RUNNER] = Field(default=RunnerOpType.UNASSIGN_RUNNER, frozen=True) + runner_id: RunnerId + +class RunnerUpOp(BaseRunnerOp[Literal[RunnerOpType.RUNNER_UP]]): + op_type: Literal[RunnerOpType.RUNNER_UP] = Field(default=RunnerOpType.RUNNER_UP, frozen=True) + runner_id: RunnerId + +class RunnerDownOp(BaseRunnerOp[Literal[RunnerOpType.RUNNER_DOWN]]): + op_type: Literal[RunnerOpType.RUNNER_DOWN] = Field(default=RunnerOpType.RUNNER_DOWN, frozen=True) + runner_id: RunnerId + +class DownloadOp(BaseRunnerOp[Literal[RunnerOpType.DOWNLOAD]]): + op_type: Literal[RunnerOpType.DOWNLOAD] = Field(default=RunnerOpType.DOWNLOAD, frozen=True) + instance_id: InstanceId + runner_id: RunnerId + shard_metadata: ShardMetadata + hosts: list[Host] + +class ExecuteTaskOp(BaseRunnerOp[Literal[RunnerOpType.CHAT_COMPLETION]]): + op_type: Literal[RunnerOpType.CHAT_COMPLETION] = Field(default=RunnerOpType.CHAT_COMPLETION, frozen=True) + runner_id: RunnerId + task: Task[TaskType, TaskStatusType] + + +# Aggregate all runner operations into a single, strictly-typed union for dispatching. +RunnerOp = Annotated[ + Union[ + AssignRunnerOp, + UnassignRunnerOp, + RunnerUpOp, + RunnerDownOp, + DownloadOp, + ExecuteTaskOp, + ], + Field(discriminator="op_type") +] \ No newline at end of file diff --git a/shared/types/worker/runners.py b/shared/types/worker/runners.py index bac23aa0..1b6c371b 100644 --- a/shared/types/worker/runners.py +++ b/shared/types/worker/runners.py @@ -1,4 +1,4 @@ -from collections.abc import Mapping, Sequence +from collections.abc import Mapping from enum import Enum from typing import Annotated, Generic, Literal, TypeVar @@ -7,71 +7,83 @@ from pydantic import BaseModel, Field, TypeAdapter, model_validator from shared.types.common import NodeId from shared.types.models.common import ModelId from shared.types.worker.common import RunnerId -from shared.types.worker.downloads import BaseDownloadProgress, DownloadStatus -from shared.types.worker.shards import PartitionStrategy, ShardMetadata +from shared.types.worker.downloads import DownloadProgress +from shared.types.worker.shards import ShardMetadata class RunnerStatusType(str, Enum): - Rejected = "Rejected" - Starting = "Starting" + Assigned = "Assigned" Downloading = "Downloading" + Ready = "Ready" + Starting = "Starting" + Loaded = "Loaded" Running = "Running" Failed = "Failed" -RunnerStatusTypeT = TypeVar("RunnerStatusTypeT", bound=RunnerStatusType) +RunnerStatusTypeT = TypeVar("RunnerStatusTypeT", bound=RunnerStatusType, covariant=True) -class RunnerStatus(BaseModel, Generic[RunnerStatusTypeT]): +class BaseRunnerStatus(BaseModel, Generic[RunnerStatusTypeT]): runner_status: RunnerStatusTypeT -class RejectedRunnerStatus(RunnerStatus[RunnerStatusType.Rejected]): - runner_status: Literal[RunnerStatusType.Rejected] +# Emitted by the Master +class AssignedRunnerStatus(BaseRunnerStatus[RunnerStatusType.Assigned]): + runner_status: Literal[RunnerStatusType.Assigned] = Field(default=RunnerStatusType.Assigned) +# Emitted by the Worker +class DownloadingRunnerStatus(BaseRunnerStatus[RunnerStatusType.Downloading]): + runner_status: Literal[RunnerStatusType.Downloading] = Field(default=RunnerStatusType.Downloading) + download_progress: DownloadProgress -class StartingRunnerStatus(RunnerStatus[RunnerStatusType.Starting]): - runner_status: Literal[RunnerStatusType.Starting] +# Emitted by the Worker +class ReadyRunnerStatus(BaseRunnerStatus[RunnerStatusType.Ready]): + runner_status: Literal[RunnerStatusType.Ready] = Field(default=RunnerStatusType.Ready) +# Emitted by the Master +class StartingRunnerStatus(BaseRunnerStatus[RunnerStatusType.Starting]): + runner_status: Literal[RunnerStatusType.Starting] = Field(default=RunnerStatusType.Starting) -class DownloadingRunnerStatus(RunnerStatus[RunnerStatusType.Downloading]): - runner_status: Literal[RunnerStatusType.Downloading] - download_progress: BaseDownloadProgress[DownloadStatus] +# Emitted by the Worker +class LoadedRunnerStatus(BaseRunnerStatus[RunnerStatusType.Loaded]): + runner_status: Literal[RunnerStatusType.Loaded] = Field(default=RunnerStatusType.Loaded) +# Emitted by the Worker +class RunningRunnerStatus(BaseRunnerStatus[RunnerStatusType.Running]): + runner_status: Literal[RunnerStatusType.Running] = Field(default=RunnerStatusType.Running) -class RunningRunnerStatus(RunnerStatus[RunnerStatusType.Running]): - runner_status: Literal[RunnerStatusType.Running] - - -class FailedRunnerStatus(RunnerStatus[RunnerStatusType.Failed]): - runner_status: Literal[RunnerStatusType.Failed] +# Emitted by the Worker +class FailedRunnerStatus(BaseRunnerStatus[RunnerStatusType.Failed]): + runner_status: Literal[RunnerStatusType.Failed] = Field(default=RunnerStatusType.Failed) error_message: str | None = None -_RunnerStatus = Annotated[ - RejectedRunnerStatus - | StartingRunnerStatus +RunnerStatus = Annotated[ + AssignedRunnerStatus | DownloadingRunnerStatus + | ReadyRunnerStatus + | StartingRunnerStatus + | LoadedRunnerStatus | RunningRunnerStatus | FailedRunnerStatus, Field, ] -RunnerStatusParser: TypeAdapter[RunnerStatus[RunnerStatusType]] = TypeAdapter( - _RunnerStatus +RunnerStatusParser: TypeAdapter[RunnerStatus] = TypeAdapter( + RunnerStatus ) class ShardAssignments(BaseModel): model_id: ModelId - runner_to_shard: Mapping[RunnerId, ShardMetadata[PartitionStrategy]] - node_to_runner: Mapping[NodeId, Sequence[RunnerId]] + runner_to_shard: Mapping[RunnerId, ShardMetadata] + node_to_runner: Mapping[NodeId, 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" - ) + for runner_id in self.node_to_runner.values(): + 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 index 67361967..5ee7baa8 100644 --- a/shared/types/worker/shards.py +++ b/shared/types/worker/shards.py @@ -1,5 +1,5 @@ from enum import Enum -from typing import Annotated, Literal +from typing import Annotated, Generic, Literal, TypeAlias, TypeVar from pydantic import BaseModel, DirectoryPath, Field, TypeAdapter @@ -11,7 +11,10 @@ class PartitionStrategy(str, Enum): pipeline = "pipeline" -class ShardMetadata[PartitionStrategyT: PartitionStrategy](BaseModel): +PartitionStrategyT = TypeVar("PartitionStrategyT", bound=PartitionStrategy, covariant=True) + + +class BaseShardMetadata(BaseModel, Generic[PartitionStrategyT]): """ Defines a specific shard of the model that is ready to be run on a device. Replaces previous `Shard` object. @@ -24,7 +27,7 @@ class ShardMetadata[PartitionStrategyT: PartitionStrategy](BaseModel): model_path: DirectoryPath -class PipelineShardMetadata(ShardMetadata[PartitionStrategy.pipeline]): +class PipelineShardMetadata(BaseShardMetadata[Literal[PartitionStrategy.pipeline]]): """ Pipeline parallelism shard meta. """ @@ -36,19 +39,30 @@ class PipelineShardMetadata(ShardMetadata[PartitionStrategy.pipeline]): end_layer: Annotated[int, Field(ge=0)] -_ShardMetadata = Annotated[ +ShardMetadata = Annotated[ PipelineShardMetadata, Field(discriminator="partition_strategy") ] -ShardMetaParser: TypeAdapter[ShardMetadata[PartitionStrategy]] = TypeAdapter( - _ShardMetadata +ShardMetadataParser: TypeAdapter[ShardMetadata] = TypeAdapter( + ShardMetadata ) +# --------------------------------------------------------------------------- +# Convenience aliases +# --------------------------------------------------------------------------- -class ShardPlacement[PartitionStrategyT: PartitionStrategy](BaseModel): +# "ShardMeta" is a widely-used alias for the concrete, fully-parameterised +# `ShardMetadata` type. Defining it here avoids repetitive generic +# parameters at call-sites and resolves unknown-import diagnostics in +# downstream modules. + +ShardMeta: TypeAlias = ShardMetadata + + +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]] + shard_assignments: dict[NodeId, BaseShardMetadata[PartitionStrategyT]] diff --git a/test_shard_serialization.py b/test_shard_serialization.py new file mode 100644 index 00000000..0519ecba --- /dev/null +++ b/test_shard_serialization.py @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/worker/main.py b/worker/main.py index e69de29b..28179437 100644 --- a/worker/main.py +++ b/worker/main.py @@ -0,0 +1,358 @@ +import asyncio +import os +from asyncio.queues import Queue +from functools import partial +from logging import Logger +from typing import AsyncGenerator, Optional + +from pydantic import BaseModel, ConfigDict + +from shared.types.common import NodeId +from shared.types.events.events import ChunkGenerated, InstanceId, RunnerStatusUpdated +from shared.types.events.registry import Event +from shared.types.states.worker import WorkerState +from shared.types.worker.common import RunnerId +from shared.types.worker.downloads import ( + DownloadCompleted, + DownloadFailed, + DownloadOngoing, + DownloadProgressData, +) +from shared.types.worker.mlx import Host +from shared.types.worker.ops import ( + AssignRunnerOp, + DownloadOp, + ExecuteTaskOp, + RunnerDownOp, + RunnerOp, + RunnerOpType, + RunnerUpOp, + UnassignRunnerOp, +) +from shared.types.worker.runners import ( + DownloadingRunnerStatus, + FailedRunnerStatus, + LoadedRunnerStatus, + ReadyRunnerStatus, + RunnerStatus, + RunnerStatusType, + RunningRunnerStatus, +) +from shared.types.worker.shards import ShardMetadata +from worker.runner.runner_supervisor import RunnerSupervisor + + +class AssignedRunner(BaseModel): + runner_id: RunnerId + instance_id: InstanceId + shard_metadata: ShardMetadata # just data + hosts: list[Host] + + status: RunnerStatus + runner: Optional[RunnerSupervisor] # set if the runner is 'up' + + model_config = ConfigDict(arbitrary_types_allowed=True) + + @property + def is_downloaded(self) -> bool: + # TODO: Do this properly with huggingface validating each of the files. + return os.path.exists(self.shard_metadata.model_path) + + def status_update_event(self) -> RunnerStatusUpdated: + return RunnerStatusUpdated( + runner_id=self.runner_id, + runner_status=self.status, + ) + +class Worker: + def __init__( + self, + node_id: NodeId, + initial_state: WorkerState, + logger: Logger, + ): + self.node_id = node_id + self.state = initial_state + self.logger = logger + + self.assigned_runners: dict[RunnerId, AssignedRunner] = {} + self._task: asyncio.Task[None] | None = None + + ## Worker lifecycle management + @property + def _is_running(self) -> bool: + return self._task is not None and not self._task.done() + + async def start(self): + self._task = asyncio.create_task(self._loop()) + + async def stop(self): + if not self._is_running: + raise RuntimeError("Worker is not running") + + assert self._task is not None + + self._task.cancel() + + ## Op Executors + + async def _execute_assign_op( + self, op: AssignRunnerOp + ) -> AsyncGenerator[Event, None]: + ''' + Here, we are sure that the model is already downloaded. + This op moves the runner from Assigned -> Ready state. + ''' + self.assigned_runners[op.runner_id] = AssignedRunner( + runner_id=op.runner_id, + instance_id=op.instance_id, + shard_metadata=op.shard_metadata, + hosts=op.hosts, + status=ReadyRunnerStatus(), + runner=None, + ) + + yield self.assigned_runners[op.runner_id].status_update_event() + + async def _execute_unassign_op( + self, op: UnassignRunnerOp + ) -> AsyncGenerator[Event, None]: + if op.runner_id not in self.assigned_runners: + return + + # We can try to do a graceful shutdown of the runner. + runner: RunnerSupervisor | None = self.assigned_runners[op.runner_id].runner + if runner is not None: + await runner.astop() + + # This is all we really need: + del self.assigned_runners[op.runner_id] + + return + yield + + async def _execute_runner_up_op( + self, op: RunnerUpOp + ) -> AsyncGenerator[Event, None]: + assigned_runner = self.assigned_runners[op.runner_id] + + assigned_runner.runner = await RunnerSupervisor.create( + model_shard_meta=assigned_runner.shard_metadata, + hosts=assigned_runner.hosts, + ) + + if assigned_runner.runner.healthy: + assigned_runner.status = LoadedRunnerStatus() + else: + assigned_runner.status = FailedRunnerStatus() + yield self.assigned_runners[op.runner_id].status_update_event() + + async def _execute_runner_down_op( + self, op: RunnerDownOp + ) -> AsyncGenerator[Event, None]: + assigned_runner = self.assigned_runners[op.runner_id] + + assert isinstance(assigned_runner.runner, RunnerSupervisor) + await assigned_runner.runner.astop() + assigned_runner.runner = None + + assigned_runner.status = ReadyRunnerStatus() + yield assigned_runner.status_update_event() + + async def _execute_download_op( + self, op: DownloadOp + ) -> AsyncGenerator[Event, None]: + ''' + The model needs assigning and then downloading. + This op moves the runner from Assigned -> Downloading -> Ready state. + ''' + initial_status = DownloadingRunnerStatus( + download_progress=DownloadOngoing( + node_id=self.node_id, + download_progress=DownloadProgressData( + total_bytes=1, # tmp + downloaded_bytes=0 + ) + ) + ) + + self.assigned_runners[op.runner_id] = AssignedRunner( + runner_id=op.runner_id, + instance_id=op.instance_id, + shard_metadata=op.shard_metadata, + hosts=op.hosts, + status=initial_status, + runner=None, + ) + assigned_runner: AssignedRunner = self.assigned_runners[op.runner_id] + yield assigned_runner.status_update_event() + + # Download it! + # TODO: we probably want download progress as part of a callback that gets passed to the downloader. + + try: + assert assigned_runner.is_downloaded + assigned_runner.status = DownloadingRunnerStatus( + download_progress=DownloadCompleted( + node_id=self.node_id, + ) + ) + except Exception as e: + assigned_runner.status = DownloadingRunnerStatus( + download_progress=DownloadFailed( + node_id=self.node_id, + error_message=str(e) + ) + ) + yield assigned_runner.status_update_event() + + assigned_runner.status = ReadyRunnerStatus() + yield assigned_runner.status_update_event() + +# Plan: +# First get a single inference running +# Then build boilerplate for passing callback when mlx is in the 'ready' state +# Then figure out if we can do what's needed with events. But this is a little challenging because it depends on Alex's code. + async def _execute_chat_completion_op( + self, op: ExecuteTaskOp + ) -> AsyncGenerator[Event, None]: + ''' + This is the entry point for a chat completion starting. + While there is only one execute function, it will get called in different ways for runner 0 and runner [1, 2, 3, ...]. + Runners [1, 2, 3, ...] will run this method when a task is in 'pending' state. + Runner 0 will run this method when a task is in 'running' state. + TODO: How do we handle the logic of ensuring that n-1 nodes have started their execution before allowing the 0'th runner to start? + This is still a little unclear to me. + ''' + assigned_runner = self.assigned_runners[op.runner_id] + + async def inner_execute(queue: asyncio.Queue[Event]) -> None: + assert assigned_runner.runner is not None + assert assigned_runner.runner.healthy + + async def running_callback(queue: asyncio.Queue[Event]) -> None: + # Called when the MLX process has been kicked off + assigned_runner.status = RunningRunnerStatus() + await queue.put(assigned_runner.status_update_event()) + + try: + async for chunk in assigned_runner.runner.stream_response( + task=op.task, + request_started_callback=partial(running_callback, queue)): + await queue.put(ChunkGenerated( + task_id=op.task.task_id, + chunk=chunk + )) + + # After a successful inference: + assigned_runner.status = LoadedRunnerStatus() + await queue.put(assigned_runner.status_update_event()) + + except Exception as e: + # TODO: What log level? + self.logger.log(2, f'Runner failed whilst running inference task. Task: {op.task}. Error: {e}') + + assigned_runner.runner = None + assigned_runner.status = FailedRunnerStatus(error_message=str(e)) + await queue.put(assigned_runner.status_update_event()) + + queue: Queue[Event] = asyncio.Queue() + task = asyncio.create_task(inner_execute(queue)) + + try: + # Yield items from the queue + while True: + item: Event = await asyncio.wait_for(queue.get(), timeout=5) + yield item + if isinstance(item, RunnerStatusUpdated) and isinstance( + item.runner_status, (LoadedRunnerStatus, FailedRunnerStatus) + ): + break + finally: + # Ensure the task is cleaned up + await task + + + ## Operation Planner + + async def _execute_op(self, op: RunnerOp) -> AsyncGenerator[Event, None]: + ## It would be great if we can get rid of this async for ... yield pattern. + match op.op_type: + case RunnerOpType.ASSIGN_RUNNER: + event_generator = self._execute_assign_op(op) + case RunnerOpType.UNASSIGN_RUNNER: + event_generator = self._execute_unassign_op(op) + case RunnerOpType.RUNNER_UP: + event_generator = self._execute_runner_up_op(op) + case RunnerOpType.RUNNER_DOWN: + event_generator = self._execute_runner_down_op(op) + case RunnerOpType.DOWNLOAD: + event_generator = self._execute_download_op(op) + case RunnerOpType.CHAT_COMPLETION: + event_generator = self._execute_chat_completion_op(op) + + async for event in event_generator: + yield event + + ## Planning logic + def plan(self, state: WorkerState) -> RunnerOp | None: + # Compare state to worker 'mood' + + # First spin things down + + # Then spin things up + + # Then make sure things are downloading. + for instance_id, instance in state.instances.instances.items(): + # We should already have asserted that this runner exists + # If it didn't exist then we return a assign_runner op. + for node_id, runner_id in instance.instance_params.shard_assignments.node_to_runner.items(): + if node_id != self.node_id: + continue + assert runner_id in self.assigned_runners + + runner = self.assigned_runners[runner_id] + + if not runner.is_downloaded: + if runner.status.runner_status == RunnerStatusType.Downloading: + return None + else: + return DownloadOp( + runner_id=runner_id, + instance_id=instance_id, + shard_metadata=instance.instance_params.shard_assignments.runner_to_shard[runner_id], + hosts=instance.instance_params.hosts + ) + + + + + # Finally, chat completion. + return None + + + # Handle state updates + async def _loop(self): + while True: + state_copy = self.state.model_copy(deep=True) + + op: RunnerOp | None = self.plan(state_copy) + + # Run the op, synchronously blocking for now. + if op is not None: + async for event in self._execute_op(op): + print(event) + # self.event_publisher(event) + + await asyncio.sleep(0.01) + + # TODO: Handle tail event log + # TODO: Handle resource monitoring (write-only) + +async def main(): + + + print("Hello from worker!") + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/worker/runner/communication.py b/worker/runner/communication.py index 5491f171..18001b8f 100644 --- a/worker/runner/communication.py +++ b/worker/runner/communication.py @@ -23,6 +23,8 @@ async def supervisor_write_message( ) encoded: bytes = message.model_dump_json().encode("utf-8") + b"\n" + print(f"message: {message}") + # print(f"encoded: {encoded}") proc.stdin.write(encoded) await proc.stdin.drain() @@ -31,7 +33,7 @@ async def runner_read_message() -> RunnerMessage: loop = asyncio.get_running_loop() line: bytes = await loop.run_in_executor(None, sys.stdin.buffer.readline) - if not line: + if not line: # This seems to be what triggers when we don't clean up the runner neatly and leave the process dangling. raise EOFError("No more data to read") line = line.strip() diff --git a/worker/runner/runner.py b/worker/runner/runner.py index 3e4d76b3..7b5b2e6d 100644 --- a/worker/runner/runner.py +++ b/worker/runner/runner.py @@ -1,8 +1,6 @@ import asyncio import concurrent.futures -from asyncio.events import AbstractEventLoop from collections.abc import AsyncGenerator -from concurrent.futures.thread import ThreadPoolExecutor from functools import partial from typing import Callable, cast @@ -13,9 +11,7 @@ from mlx_lm.tokenizer_utils import TokenizerWrapper from engines.mlx.utils_mlx import apply_chat_template, initialize_mlx from shared.openai import FinishReason -from shared.types.tasks.common import ( - TaskData, -) +from shared.types.tasks.common import ChatCompletionTaskData, CompletionCreateParams from shared.types.worker.commands_runner import ( ChatTaskMessage, ExitMessage, @@ -24,8 +20,6 @@ from shared.types.worker.commands_runner import ( RunnerMessage, SetupMessage, ) -from shared.types.worker.mlx import Host -from shared.types.worker.shards import ShardMeta from shared.utils import ensure_type from worker.runner.communication import ( runner_print, @@ -40,7 +34,7 @@ async def _mlx_generate( model: nn.Module, tokenizer: TokenizerWrapper, sampler: Callable[[mx.array], mx.array], - task: TaskData, + task: ChatCompletionTaskData, ) -> AsyncGenerator[GenerationResponse]: loop = asyncio.get_running_loop() queue: asyncio.Queue[GenerationResponse | Exception | object] = asyncio.Queue() @@ -69,17 +63,17 @@ async def _mlx_generate( _ = loop.call_soon_threadsafe(queue.put_nowait, sentinel) # Currently we support chat-completion tasks only. - task_data = task.task_data + task_data: CompletionCreateParams = task.task_params runner_print(f"task_data: {task_data}") prompt = await apply_chat_template( mlx_executor=mlx_executor, tokenizer=tokenizer, - chat_task=task_data, + chat_task_data=task_data, ) - max_tokens = task_data.max_tokens or 100 + max_tokens = task.task_params.max_tokens or 100 generation_fn = partial(_generate_tokens, prompt, max_tokens) future = loop.run_in_executor(mlx_executor, generation_fn) @@ -94,9 +88,12 @@ async def _mlx_generate( if isinstance(item, Exception): raise item + assert isinstance(item, GenerationResponse) # constrain datatype + runner_print(item.text) yield item + # TODO: There is a big bug on this line! assert future.done() @@ -105,17 +102,15 @@ async def main(): runner_print("hello from the runner") # Get setup info from worker - init_message: RunnerMessage = await runner_read_message() - setup_message: SetupMessage = ensure_type(init_message, SetupMessage) - model_shard_meta: ShardMeta = setup_message.model_shard_meta - hosts: list[Host] = setup_message.hosts + init_message = await runner_read_message() + setup_message = ensure_type(init_message, SetupMessage) + model_shard_meta = setup_message.model_shard_meta + hosts = setup_message.hosts - mlx_executor: ThreadPoolExecutor = concurrent.futures.ThreadPoolExecutor( - max_workers=1 - ) - loop: AbstractEventLoop = asyncio.get_running_loop() + mlx_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) + loop = asyncio.get_running_loop() - runner_print(f"got here; {model_shard_meta.model_path}") + runner_print(f"got here; {hosts}") model, tokenizer, sampler = await loop.run_in_executor( mlx_executor, @@ -125,13 +120,12 @@ async def main(): while True: message: RunnerMessage = await runner_read_message() match message: - case ChatTaskMessage(task=task_data): + case ChatTaskMessage(task_data=task_data): runner_print(f"received chat request: {task_data}") - # Ensure we have a chat-completion task subtype - messages = task_data.task_data.messages - messages_dicts = [msg.model_dump() for msg in messages] - runner_print(f"messages_dicts RUNNER: {messages_dicts}") + prompt = task_data.task_params.messages[0] + if prompt.content is not None and 'EXO RUNNER MUST FAIL' in prompt.content: + raise Exception('Artificial runner exception - for testing purposes only.') # Generate responses using the actual MLX generation async for generation_response in _mlx_generate( diff --git a/worker/runner/runner_supervisor.py b/worker/runner/runner_supervisor.py index ba15bf4a..5ca77bfc 100644 --- a/worker/runner/runner_supervisor.py +++ b/worker/runner/runner_supervisor.py @@ -2,10 +2,16 @@ import asyncio import contextlib import sys from collections.abc import AsyncGenerator +from types import CoroutineType from typing import Any, Callable from shared.types.events.chunks import GenerationChunk, TokenChunk, TokenChunkData -from shared.types.tasks.common import Task, TaskStatusType, TaskType +from shared.types.tasks.common import ( + ChatCompletionTaskData, + Task, + TaskStatusTypeT, + TaskTypeT, +) from shared.types.worker.commands_runner import ( ChatTaskMessage, ErrorResponse, @@ -31,14 +37,15 @@ class RunnerSupervisor: Use the class method `create` to properly initialize an instance. """ + # TODO: Logger. def __init__( self, - model_shard_meta: ShardMetadata[Any], + model_shard_meta: ShardMetadata, hosts: list[Host], runner_process: asyncio.subprocess.Process, ): """Private constructor. Use RunnerSupervisor.create() instead.""" - self.model_shard_meta: ShardMetadata[Any] = model_shard_meta + self.model_shard_meta: ShardMetadata = model_shard_meta self.hosts: list[Host] = hosts self.runner_process: asyncio.subprocess.Process = runner_process self.running: bool = True @@ -50,7 +57,7 @@ class RunnerSupervisor: @classmethod async def create( cls, - model_shard_meta: ShardMetadata[Any], + model_shard_meta: ShardMetadata, hosts: list[Host], ) -> "RunnerSupervisor": """ @@ -68,6 +75,7 @@ class RunnerSupervisor: ) ) + print(f'{model_shard_meta=}') await supervisor_write_message( runner_process, SetupMessage( @@ -140,8 +148,8 @@ class RunnerSupervisor: async def stream_response( self, - task: Task[TaskType, TaskStatusType], - request_started_callback: Callable[[], None] | None = None, + task: Task[TaskTypeT, TaskStatusTypeT], + request_started_callback: Callable[..., CoroutineType[Any, Any, None]] | None = None, # fyi this is async now ) -> AsyncGenerator[GenerationChunk]: """ Streams a chat request from the model. @@ -151,13 +159,20 @@ class RunnerSupervisor: if not self.healthy: raise RuntimeError("Runner process was found to be dead") + task_data = task.task_data + assert isinstance(task_data, ChatCompletionTaskData) # this is messy for now. await supervisor_write_message( proc=self.runner_process, message=ChatTaskMessage( - task=task.task_data, + task_data=task_data, ), ) + # This is easy for now. If we need more reliability, the runner can have a new 'ready' message type. + if request_started_callback is not None: + await request_started_callback() + + while True: line: RunnerResponse | None = await supervisor_read_response( self.runner_process diff --git a/worker/test_worker_state.py b/worker/test_worker_state.py new file mode 100644 index 00000000..5db3f9a9 --- /dev/null +++ b/worker/test_worker_state.py @@ -0,0 +1,48 @@ +## Tests for worker state differentials +## When the worker state changes, this should be reflected by a worker intention. + +import asyncio +from typing import Callable +from uuid import uuid4 + +import pytest + +from shared.types.common import NodeId +from shared.types.states.worker import NodeStatusState, WorkerState +from shared.types.worker.common import InstanceId, NodeStatus +from shared.types.worker.instances import Instance +from worker.main import Worker + + +@pytest.mark.asyncio +async def test_worker_runs_and_stops(worker: Worker): + await worker.start() + await asyncio.sleep(0.01) + + assert worker._is_running # type: ignore + + await worker.stop() + await asyncio.sleep(0.01) + + assert not worker._is_running # type: ignore + +@pytest.mark.asyncio +async def test_worker_instance_added(worker: Worker, instance: Callable[[NodeId], Instance]): + await worker.start() + await asyncio.sleep(0.01) + + worker.state.instances.instances = {InstanceId(uuid4()): instance(worker.node_id)} + + print(worker.state.instances.instances) + +def test_plan_noop(worker: Worker): + s = WorkerState( + node_status=NodeStatusState( + node_status={ + NodeId(uuid4()): NodeStatus.Idle + } + ), + ) + next_op = worker.plan(s) + + assert next_op is None diff --git a/worker/tests/conftest.py b/worker/tests/conftest.py index a631cb4c..c8687a04 100644 --- a/worker/tests/conftest.py +++ b/worker/tests/conftest.py @@ -1,39 +1,39 @@ +import asyncio import uuid +from logging import Logger, getLogger from pathlib import Path -from typing import Callable, cast +from typing import Callable, Literal import pytest -from openai.types.chat import ChatCompletionUserMessageParam -from openai.types.chat.completion_create_params import ( - CompletionCreateParamsNonStreaming, - CompletionCreateParamsStreaming, -) -from pydantic import TypeAdapter +from shared.types.common import NodeId from shared.types.models.common import ModelId +from shared.types.states.worker import NodeStatusState, WorkerState from shared.types.tasks.common import ( - ChatCompletionStreamingTask, + ChatCompletionMessage, + ChatCompletionTaskData, + CompletionCreateParams, Task, TaskArtifact, TaskId, TaskState, TaskStatusOtherType, - TaskStatusType, TaskType, ) -from shared.types.worker.common import InstanceId +from shared.types.worker.common import InstanceId, NodeStatus +from shared.types.worker.instances import Instance, InstanceParams, TypeOfInstance from shared.types.worker.mlx import Host -from shared.types.worker.shards import PipelineShardMetadata - -CompletionCreateParamsStreamingAdapter = TypeAdapter(CompletionCreateParamsStreaming) -CompletionCreateParamsNonStreamingAdapter = TypeAdapter( - CompletionCreateParamsNonStreaming +from shared.types.worker.ops import ( + AssignRunnerOp, + RunnerUpOp, ) +from shared.types.worker.runners import RunnerId, ShardAssignments +from shared.types.worker.shards import PipelineShardMetadata +from worker.main import Worker -# Concrete TaskArtifact implementation for pending streaming tasks class PendingStreamingTaskArtifact( - TaskArtifact[TaskType.ChatCompletionStreaming, TaskStatusOtherType.Pending] + TaskArtifact[Literal[TaskType.ChatCompletion], Literal[TaskStatusOtherType.Pending]] ): pass @@ -97,38 +97,119 @@ def user_message(): @pytest.fixture -def chat_completion_params(user_message: str): +def completion_create_params(user_message: str) -> CompletionCreateParams: """Creates ChatCompletionParams with the given message""" - return CompletionCreateParamsStreaming( + return CompletionCreateParams( model="gpt-4", - messages=[ChatCompletionUserMessageParam(role="user", content=user_message)], + messages=[ChatCompletionMessage(role="user", content=user_message)], stream=True, ) +@pytest.fixture +def chat_completion_task(completion_create_params: CompletionCreateParams) -> ChatCompletionTaskData: + """Creates a ChatCompletionTask directly for serdes testing""" + return ChatCompletionTaskData(task_params=completion_create_params) @pytest.fixture -def chat_completion_streaming_task_data( - chat_completion_params: CompletionCreateParamsStreaming, -): - """Creates ChatCompletionStreamingTask from params""" - return ChatCompletionStreamingTask(task_data=chat_completion_params) - - -@pytest.fixture -def streaming_task( - chat_completion_streaming_task_data: CompletionCreateParamsStreaming, -) -> Task[TaskType, TaskStatusType]: +def chat_task( + completion_create_params: CompletionCreateParams, +) -> Task[Literal[TaskType.ChatCompletion], TaskStatusOtherType]: """Creates the final Task object""" - task = Task( + return Task[Literal[TaskType.ChatCompletion], TaskStatusOtherType]( task_id=TaskId(), - task_type=TaskType.ChatCompletionStreaming, - task_params=ChatCompletionStreamingTask( - task_data=chat_completion_streaming_task_data + task_type=TaskType.ChatCompletion, + task_data=ChatCompletionTaskData( + task_params=completion_create_params ), - task_state=TaskState( + task_state=TaskState[TaskStatusOtherType, Literal[TaskType.ChatCompletion]]( task_status=TaskStatusOtherType.Pending, task_artifact=PendingStreamingTaskArtifact(), ), on_instance=InstanceId(), ) - return cast(Task[TaskType, TaskStatusType], task) + +@pytest.fixture +def worker_state(): + node_status=NodeStatusState( + node_status={ + NodeId(uuid.uuid4()): NodeStatus.Idle + } + ) + + return WorkerState( + node_status=node_status, + ) + +@pytest.fixture +def logger() -> Logger: + return getLogger("test_logger") + +@pytest.fixture +def instance(pipeline_shard_meta: Callable[[int, int], PipelineShardMetadata], hosts_one: list[Host]): + def _instance(node_id: NodeId) -> Instance: + model_id = ModelId(uuid.uuid4()) + runner_id = RunnerId(uuid.uuid4()) + + shard_assignments = ShardAssignments( + model_id=model_id, + runner_to_shard={ + runner_id: pipeline_shard_meta(1, 0) + }, + node_to_runner={node_id: runner_id} + ) + + instance_params = InstanceParams( + shard_assignments=shard_assignments, + hosts=hosts_one + ) + + return Instance( + instance_id=InstanceId(uuid.uuid4()), + instance_params=instance_params, + instance_type=TypeOfInstance.ACTIVE + ) + return _instance + +@pytest.fixture +def worker(worker_state: WorkerState, logger: Logger): + return Worker(NodeId(uuid.uuid4()), worker_state, logger) + +@pytest.fixture +async def worker_with_assigned_runner(worker: Worker, instance: Callable[[NodeId], Instance]): + """Fixture that provides a worker with an already assigned runner.""" + await worker.start() + await asyncio.sleep(0.01) + + instance_obj: Instance = instance(worker.node_id) + + # Extract runner_id from shard assignments + runner_id = next(iter(instance_obj.instance_params.shard_assignments.runner_to_shard)) + + # Assign the runner + assign_op = AssignRunnerOp( + runner_id=runner_id, + shard_metadata=instance_obj.instance_params.shard_assignments.runner_to_shard[runner_id], + hosts=instance_obj.instance_params.hosts, + instance_id=instance_obj.instance_id, + ) + + async for _ in worker._execute_op(assign_op): # type: ignore[misc] + pass + + return worker, runner_id, instance_obj + +@pytest.fixture +async def worker_with_running_runner(worker_with_assigned_runner: tuple[Worker, RunnerId, Instance]): + """Fixture that provides a worker with an already assigned runner.""" + worker, runner_id, instance_obj = worker_with_assigned_runner + + runner_up_op = RunnerUpOp(runner_id=runner_id) + async for _ in worker._execute_op(runner_up_op): # type: ignore[misc] + pass + + # Is the runner actually running? + supervisor = next(iter(worker.assigned_runners.values())).runner + assert supervisor is not None + assert supervisor.healthy + + return worker, runner_id, instance_obj \ No newline at end of file diff --git a/worker/tests/test_serdes.py b/worker/tests/test_serdes.py index 8119aa4a..187c4dfd 100644 --- a/worker/tests/test_serdes.py +++ b/worker/tests/test_serdes.py @@ -1,8 +1,8 @@ -from typing import Callable, Literal, TypeVar +from typing import Callable, TypeVar from pydantic import BaseModel, TypeAdapter -from shared.types.tasks.common import Task, TaskStatusOtherType, TaskType +from shared.types.tasks.common import ChatCompletionTaskData from shared.types.worker.commands_runner import ( ChatTaskMessage, RunnerMessageTypeAdapter, @@ -35,9 +35,9 @@ def test_supervisor_setup_message_serdes( def test_supervisor_task_message_serdes( - streaming_task: Task[TaskType, Literal[TaskStatusOtherType.Pending]], + chat_completion_task: ChatCompletionTaskData, ): task_message = ChatTaskMessage( - task=streaming_task.task_data, + task_data=chat_completion_task, ) assert_equal_serdes(task_message, RunnerMessageTypeAdapter) diff --git a/worker/tests/test_supervisor.py b/worker/tests/test_supervisor.py index 3c17099d..b63233be 100644 --- a/worker/tests/test_supervisor.py +++ b/worker/tests/test_supervisor.py @@ -1,13 +1,19 @@ import asyncio -from typing import Callable +from typing import Callable, Literal import pytest from shared.openai import FinishReason from shared.types.events.chunks import TokenChunk -from shared.types.tasks.common import Task, TaskStatusType, TaskType +from shared.types.tasks.common import ( + ChatCompletionTaskData, + Task, + TaskStatusOtherType, + TaskStatusType, + TaskType, +) from shared.types.worker.mlx import Host -from shared.types.worker.shards import PipelineShardMeta +from shared.types.worker.shards import PipelineShardMetadata from worker.runner.runner_supervisor import RunnerSupervisor @@ -19,13 +25,15 @@ def user_message(): @pytest.mark.asyncio async def test_supervisor_single_node_response( - pipeline_shard_meta: Callable[..., PipelineShardMeta], + pipeline_shard_meta: Callable[..., PipelineShardMetadata], hosts: Callable[..., list[Host]], - streaming_task: Task[TaskType, TaskStatusType], + chat_task: Task[TaskType, TaskStatusType], ): """Test that asking for the capital of France returns 'Paris' in the response""" model_shard_meta = pipeline_shard_meta(1, 0) + print(f'{model_shard_meta=}') + supervisor = await RunnerSupervisor.create( model_shard_meta=model_shard_meta, hosts=hosts(1, offset=10), @@ -35,7 +43,7 @@ async def test_supervisor_single_node_response( full_response = "" stop_reason: FinishReason | None = None - async for chunk in supervisor.stream_response(task=streaming_task): + async for chunk in supervisor.stream_response(task=chat_task): if isinstance(chunk, TokenChunk): full_response += chunk.chunk_data.text if chunk.chunk_data.finish_reason: @@ -53,9 +61,9 @@ async def test_supervisor_single_node_response( @pytest.mark.asyncio async def test_supervisor_two_node_response( - pipeline_shard_meta: Callable[..., PipelineShardMeta], + pipeline_shard_meta: Callable[..., PipelineShardMetadata], hosts: Callable[..., list[Host]], - streaming_task: Task[TaskType, TaskStatusType], + chat_task: Task[TaskType, TaskStatusType], ): """Test that asking for the capital of France returns 'Paris' in the response""" supervisor_0 = await RunnerSupervisor.create( @@ -76,13 +84,13 @@ async def test_supervisor_two_node_response( async def collect_response_0(): nonlocal full_response_0 - async for chunk in supervisor_0.stream_response(task=streaming_task): + async for chunk in supervisor_0.stream_response(task=chat_task): if isinstance(chunk, TokenChunk): full_response_0 += chunk.chunk_data.text async def collect_response_1(): nonlocal full_response_1 - async for chunk in supervisor_1.stream_response(task=streaming_task): + async for chunk in supervisor_1.stream_response(task=chat_task): if isinstance(chunk, TokenChunk): full_response_1 += chunk.chunk_data.text @@ -107,9 +115,9 @@ async def test_supervisor_two_node_response( @pytest.mark.asyncio async def test_supervisor_early_stopping( - pipeline_shard_meta: Callable[..., PipelineShardMeta], + pipeline_shard_meta: Callable[..., PipelineShardMetadata], hosts: Callable[..., list[Host]], - streaming_task: Task[TaskType, TaskStatusType], + chat_task: Task[Literal[TaskType.ChatCompletion], TaskStatusOtherType], ): """Test that asking for the capital of France returns 'Paris' in the response""" model_shard_meta = pipeline_shard_meta(1, 0) @@ -120,18 +128,23 @@ async def test_supervisor_early_stopping( ) max_tokens = 50 + assert chat_task.task_type == TaskType.ChatCompletion + print(f'chat_task.task_data: {type(chat_task.task_data)}') + assert isinstance(chat_task.task_data, ChatCompletionTaskData) + task_data: ChatCompletionTaskData = chat_task.task_data try: - streaming_task.task_data.task_data.max_tokens = max_tokens - streaming_task.task_data.task_data.messages[ - 0 - ].content = "Please count from 1 to 100" + task_data.task_params.max_tokens = max_tokens + # Convert messages to a list to allow indexing, then update the first message's content + messages = list(task_data.task_params.messages) + messages[0].content = "Please count from 1 to 100" + task_data.task_params.messages = messages full_response = "" count = 0 stop_reason: FinishReason | None = None - async for chunk in supervisor.stream_response(task=streaming_task): + async for chunk in supervisor.stream_response(task=chat_task): if isinstance(chunk, TokenChunk): full_response += chunk.chunk_data.text count += 1 @@ -152,9 +165,9 @@ async def test_supervisor_early_stopping( @pytest.mark.asyncio async def test_supervisor_handles_terminated_runner( - pipeline_shard_meta: Callable[..., PipelineShardMeta], + pipeline_shard_meta: Callable[..., PipelineShardMetadata], hosts: Callable[..., list[Host]], - streaming_task: Task[TaskType, TaskStatusType], + chat_task: Task[TaskType, TaskStatusType], ): """Test that the supervisor handles a terminated runner""" model_shard_meta = pipeline_shard_meta(1, 0) @@ -176,9 +189,9 @@ async def test_supervisor_handles_terminated_runner( @pytest.mark.asyncio async def test_supervisor_handles_killed_runner( - pipeline_shard_meta: Callable[..., PipelineShardMeta], + pipeline_shard_meta: Callable[..., PipelineShardMetadata], hosts: Callable[..., list[Host]], - streaming_task: Task[TaskType, TaskStatusType], + chat_task: Task[TaskType, TaskStatusType], ): """Test that the supervisor handles a killed runner""" model_shard_meta = pipeline_shard_meta(1, 0) diff --git a/worker/tests/test_worker_handlers.py b/worker/tests/test_worker_handlers.py new file mode 100644 index 00000000..d542664d --- /dev/null +++ b/worker/tests/test_worker_handlers.py @@ -0,0 +1,211 @@ +## Tests for worker state handlers + +import asyncio +from typing import Callable + +import pytest + +from shared.types.common import NodeId +from shared.types.events.chunks import TokenChunk, TokenChunkData +from shared.types.events.events import ChunkGenerated, RunnerStatusUpdated +from shared.types.events.registry import Event +from shared.types.tasks.common import Task, TaskStatusType, TaskType +from shared.types.worker.common import RunnerId +from shared.types.worker.instances import Instance +from shared.types.worker.ops import ( + AssignRunnerOp, + DownloadOp, + ExecuteTaskOp, + RunnerDownOp, + RunnerUpOp, + UnassignRunnerOp, +) +from shared.types.worker.runners import ( + FailedRunnerStatus, + LoadedRunnerStatus, + ReadyRunnerStatus, + RunningRunnerStatus, +) +from worker.main import Worker + + +@pytest.fixture +def user_message(): + """Override the default message to ask about France's capital""" + return "What, according to Douglas Adams, is the meaning of life, the universe and everything?" + +@pytest.mark.asyncio +async def test_assign_op(worker: Worker, instance: Callable[[NodeId], Instance]): + await worker.start() + await asyncio.sleep(0.01) + + instance_obj: Instance = instance(worker.node_id) + runner_id: RunnerId | None = None + for x in instance_obj.instance_params.shard_assignments.runner_to_shard: + runner_id = x + assert runner_id is not None + + assign_op = AssignRunnerOp( + runner_id=runner_id, + shard_metadata=instance_obj.instance_params.shard_assignments.runner_to_shard[runner_id], + hosts=instance_obj.instance_params.hosts, + instance_id=instance_obj.instance_id, + ) + + events: list[Event] = [] + + async for event in worker._execute_op(assign_op): # type: ignore[misc] + events.append(event) + + # We should have a status update saying 'starting'. + assert len(events) == 1 + assert isinstance(events[0], RunnerStatusUpdated) + assert isinstance(events[0].runner_status, ReadyRunnerStatus) + + # And the runner should be assigned + assert runner_id in worker.assigned_runners + assert isinstance(worker.assigned_runners[runner_id].status, ReadyRunnerStatus) + +@pytest.mark.asyncio +async def test_unassign_op(worker_with_assigned_runner: tuple[Worker, RunnerId, Instance]): + worker, runner_id, _ = worker_with_assigned_runner + + unassign_op = UnassignRunnerOp( + runner_id=runner_id + ) + + events: list[Event] = [] + + async for event in worker._execute_op(unassign_op): # type: ignore[misc] + events.append(event) + + # We should have no assigned runners and no events were emitted + assert len(worker.assigned_runners) == 0 + assert len(events) == 0 + +@pytest.mark.asyncio +async def test_runner_up_op(worker_with_assigned_runner: tuple[Worker, RunnerId, Instance], chat_task: Task[TaskType, TaskStatusType]): + worker, runner_id, _ = worker_with_assigned_runner + + runner_up_op = RunnerUpOp(runner_id=runner_id) + + events: list[Event] = [] + async for event in worker._execute_op(runner_up_op): # type: ignore[misc] + events.append(event) + + assert len(events) == 1 + assert isinstance(events[0], RunnerStatusUpdated) + assert isinstance(events[0].runner_status, LoadedRunnerStatus) + + # Is the runner actually running? + supervisor = next(iter(worker.assigned_runners.values())).runner + assert supervisor is not None + assert supervisor.healthy + + full_response = '' + + async for chunk in supervisor.stream_response(task=chat_task): + if isinstance(chunk, TokenChunk): + full_response += chunk.chunk_data.text + + assert "42" in full_response.lower(), ( + f"Expected '42' in response, but got: {full_response}" + ) + + runner = worker.assigned_runners[runner_id].runner + assert runner is not None + await runner.astop() # Neat cleanup. + +@pytest.mark.asyncio +async def test_runner_down_op(worker_with_running_runner: tuple[Worker, RunnerId, Instance]): + worker, runner_id, _ = worker_with_running_runner + + runner_down_op = RunnerDownOp(runner_id=runner_id) + events: list[Event] = [] + async for event in worker._execute_op(runner_down_op): # type: ignore[misc] + events.append(event) + + assert len(events) == 1 + assert isinstance(events[0], RunnerStatusUpdated) + assert isinstance(events[0].runner_status, ReadyRunnerStatus) + +@pytest.mark.asyncio +async def test_download_op(worker_with_assigned_runner: tuple[Worker, RunnerId, Instance]): + worker, runner_id, instance_obj = worker_with_assigned_runner + + print(f'{worker.assigned_runners=}') + + download_op = DownloadOp( + instance_id=instance_obj.instance_id, + runner_id=runner_id, + shard_metadata=instance_obj.instance_params.shard_assignments.runner_to_shard[runner_id], + hosts=instance_obj.instance_params.hosts, + ) + + events: list[Event] = [] + + async for event in worker._execute_op(download_op): # type: ignore[misc] + events.append(event) + + # Should give download status and then a final download status with DownloadCompleted + print(events) + +@pytest.mark.asyncio +async def test_execute_task_op( + worker_with_running_runner: tuple[Worker, RunnerId, Instance], + chat_task: Task[TaskType, TaskStatusType]): + worker, runner_id, _ = worker_with_running_runner + + execute_task_op = ExecuteTaskOp( + runner_id=runner_id, + task=chat_task + ) + + events: list[Event] = [] + async for event in worker._execute_op(execute_task_op): # type: ignore[misc] + events.append(event) + + assert len(events) > 20 + + assert isinstance(events[0], RunnerStatusUpdated) + assert isinstance(events[0].runner_status, RunningRunnerStatus) + + assert isinstance(events[-1], RunnerStatusUpdated) + assert isinstance(events[-1].runner_status, LoadedRunnerStatus) # It should not have failed. + + gen_events: list[ChunkGenerated] = [x for x in events if isinstance(x, ChunkGenerated)] + text_chunks: list[TokenChunkData] = [x.chunk.chunk_data for x in gen_events if isinstance(x.chunk.chunk_data, TokenChunkData)] + assert len(text_chunks) == len(events) - 2 + + output_text = ''.join([x.text for x in text_chunks]) + assert '42' in output_text + + runner = worker.assigned_runners[runner_id].runner + assert runner is not None + await runner.astop() # Neat cleanup. + +@pytest.mark.asyncio +async def test_execute_task_fails( + worker_with_running_runner: tuple[Worker, RunnerId, Instance], + chat_task: Task[TaskType, TaskStatusType]): + worker, runner_id, _ = worker_with_running_runner + + messages = chat_task.task_data.task_params.messages + messages[0].content = 'Artificial prompt: EXO RUNNER MUST FAIL' + + execute_task_op = ExecuteTaskOp( + runner_id=runner_id, + task=chat_task + ) + + events: list[Event] = [] + async for event in worker._execute_op(execute_task_op): # type: ignore[misc] + events.append(event) + + assert len(events) == 2 + + assert isinstance(events[0], RunnerStatusUpdated) + assert isinstance(events[0].runner_status, RunningRunnerStatus) # It tried to start. + + assert isinstance(events[-1], RunnerStatusUpdated) + assert isinstance(events[-1].runner_status, FailedRunnerStatus) # It should have failed. \ No newline at end of file diff --git a/worker/tests/test_worker_plan.py b/worker/tests/test_worker_plan.py new file mode 100644 index 00000000..02603b85 --- /dev/null +++ b/worker/tests/test_worker_plan.py @@ -0,0 +1,263 @@ +from __future__ import annotations + +import logging +from dataclasses import dataclass +from pathlib import Path +from typing import Final, List, Optional, Type + +import pytest + +from shared.types.common import NodeId +from shared.types.models.common import ModelId +from shared.types.states.shared import Instances + +# WorkerState import below after RunnerCase definition to avoid forward reference issues +from shared.types.states.worker import NodeStatusState, WorkerState +from shared.types.worker.common import InstanceId, NodeStatus, RunnerId +from shared.types.worker.downloads import DownloadOngoing, DownloadProgressData +from shared.types.worker.instances import Instance, InstanceParams, TypeOfInstance +from shared.types.worker.ops import DownloadOp +from shared.types.worker.runners import ( + DownloadingRunnerStatus, + ReadyRunnerStatus, + RunnerStatus, + ShardAssignments, +) +from shared.types.worker.shards import PipelineShardMetadata +from worker.main import AssignedRunner, Worker + + +@dataclass(slots=True, frozen=True) +class RunnerCase: + """Important, minimal state for a *single* runner relevant to planning.""" + + status: RunnerStatus + downloaded: bool # Does the model shard already exist on disk? + + +@dataclass(slots=True, frozen=True) +class PlanTestCase: + """Table-driven description of an entire planning scenario.""" + + description: str + runners: List[RunnerCase] + # If we expect an op, specify the precise type and the index of the runner it targets. + expected_op_type: Optional[Type[DownloadOp]] # Currently only DownloadOp handled. + expected_op_runner_idx: Optional[int] = None + # Allow overriding the WorkerState passed to Worker.plan. When None, a default state + # is constructed from `runners` via helper `_build_worker_state`. + worker_state_override: Optional[WorkerState] = None + + def id(self) -> str: # noqa: D401 + return self.description.replace(" ", "_") + + +def _make_downloading_status(node_id: NodeId) -> DownloadingRunnerStatus: + """Factory for a *Downloading* status with placeholder progress.""" + return DownloadingRunnerStatus( + download_progress=DownloadOngoing( + node_id=node_id, + download_progress=DownloadProgressData(total_bytes=1, downloaded_bytes=0), + ) + ) + + +# --------------------------------------------------------------------------- +# Scenarios +# --------------------------------------------------------------------------- + +TEST_CASES: Final[List[PlanTestCase]] = [ + PlanTestCase( + description="no runners ⇢ no-op", + runners=[], + expected_op_type=None, + expected_op_runner_idx=None, + ), + PlanTestCase( + description="single ready runner, model missing ⇢ expect DownloadOp", + runners=[ + RunnerCase(status=ReadyRunnerStatus(), downloaded=False), + ], + expected_op_type=DownloadOp, + expected_op_runner_idx=0, + ), + PlanTestCase( + description="runner already downloading ⇢ no-op", + runners=[ + RunnerCase(status=_make_downloading_status(NodeId()), downloaded=False), + ], + expected_op_type=None, + expected_op_runner_idx=None, + ), + PlanTestCase( + description="ready runner, model present ⇢ no-op", + runners=[ + RunnerCase(status=ReadyRunnerStatus(), downloaded=True), + ], + expected_op_type=None, + expected_op_runner_idx=None, + ), + PlanTestCase( + description="instance for other node ⇢ no-op", + runners=[ + RunnerCase(status=ReadyRunnerStatus(), downloaded=False), + ], + expected_op_type=None, + expected_op_runner_idx=None, + worker_state_override=WorkerState( + node_status=NodeStatusState(node_status={NodeId(): NodeStatus.Idle}), + instances=Instances(instances={}), + ), + ), +] + + +# --------------------------------------------------------------------------- +# Shared factory helpers +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True, slots=True) +class RunnerContext: + runner_id: RunnerId + instance_id: InstanceId + shard_metadata: PipelineShardMetadata + instance_params: InstanceParams + + +def _build_worker_state( + *, + tmp_path: Path, + node_id: NodeId, + runner_cases: List[RunnerCase], +) -> tuple[WorkerState, List[RunnerContext]]: + """Construct a WorkerState plus per-runner context objects.""" + + instances: dict[InstanceId, Instance] = {} + runner_contexts: list[RunnerContext] = [] + + for idx, _ in enumerate(runner_cases): + runner_id = RunnerId() + instance_id = InstanceId() + model_id = ModelId() + + # Unique sub-directory per runner to allow selective `downloaded` mocking. + model_subdir = tmp_path / f"runner_{idx}" + model_subdir.mkdir(exist_ok=True) + + shard_metadata = PipelineShardMetadata( + device_rank=0, + world_size=1, + model_id=model_id, + model_path=model_subdir, + start_layer=0, + end_layer=0, + ) + + shard_assignments = ShardAssignments( + model_id=model_id, + runner_to_shard={runner_id: shard_metadata}, + node_to_runner={node_id: runner_id}, + ) + + instance_params = InstanceParams( + shard_assignments=shard_assignments, + hosts=[], + ) + + instance = Instance( + instance_id=instance_id, + instance_params=instance_params, + instance_type=TypeOfInstance.ACTIVE, + ) + + instances[instance_id] = instance + + runner_contexts.append( + RunnerContext( + runner_id=runner_id, + instance_id=instance_id, + shard_metadata=shard_metadata, + instance_params=instance_params, + ) + ) + + worker_state = WorkerState( + node_status=NodeStatusState(node_status={node_id: NodeStatus.Idle}), + instances=Instances(instances=instances), + ) + + return worker_state, runner_contexts + + +# --------------------------------------------------------------------------- +# Parametrised test +# --------------------------------------------------------------------------- + + +# Pre-compute readable identifiers for each case to avoid lambda typing issues. +@pytest.mark.parametrize("case", TEST_CASES, ids=[case.id() for case in TEST_CASES]) +def test_worker_plan(case: PlanTestCase, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """Exercise Worker.plan across declarative scenarios.""" + + # Fresh identifier for isolation of node + node_id = NodeId() + + # Assemble WorkerState and surrounding objects --------------------------------------- + worker_state, runner_contexts = _build_worker_state( + tmp_path=tmp_path, + node_id=node_id, + runner_cases=case.runners, + ) + + # Replace with explicit override if provided by the scenario. + if case.worker_state_override is not None: + worker_state = case.worker_state_override + + logger = logging.getLogger("test_worker_plan") + worker = Worker(node_id=node_id, initial_state=worker_state, logger=logger) + + # Build assigned_runners and a path→downloaded lookup -------------------------------- + path_downloaded_map: dict[str, bool] = {} + + for idx, runner_case in enumerate(case.runners): + runner_status = runner_case.status + ctx = runner_contexts[idx] + + assigned_runner = AssignedRunner( + runner_id=ctx.runner_id, + instance_id=ctx.instance_id, + shard_metadata=ctx.shard_metadata, + hosts=ctx.instance_params.hosts, + status=runner_status, + runner=None, + ) + worker.assigned_runners[ctx.runner_id] = assigned_runner + + path_downloaded_map[str(ctx.shard_metadata.model_path)] = runner_case.downloaded + + # Stub filesystem existence check ------------------------------------------------------ + from worker import main as worker_main # local import for module-scoped os + + def _fake_exists(path: str | Path) -> bool: # noqa: ANN001 – match os.path.exists signature + return path_downloaded_map.get(str(path), False) + + monkeypatch.setattr(worker_main.os.path, "exists", _fake_exists) + + # Plan and assert ---------------------------------------------------------------------- + op = worker.plan(worker_state) + + if case.expected_op_type is None: + assert op is None, f"Unexpected op {op} for scenario: {case.description}" + else: + assert isinstance(op, case.expected_op_type), ( + f"Expected {case.expected_op_type.__name__}, got {type(op).__name__ if op else 'None'}" + ) + + assert case.expected_op_runner_idx is not None, "Runner index must be set when expecting an op" + target_ctx = runner_contexts[case.expected_op_runner_idx] + + assert op.runner_id == target_ctx.runner_id + assert op.instance_id == target_ctx.instance_id + assert op.shard_metadata == target_ctx.shard_metadata +