New worker

Co-authored-by: Matt Beton <matthew.beton@gmail.com>

.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

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

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

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):

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: ...
 

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
 

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
     )

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):

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]
 
 

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

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]):

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: ...

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):

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

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")
+]

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

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]]

test_shard_serialization.py

@@ -0,0 +1 @@
+ 

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())

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()
 

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(

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

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

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

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)

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)

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. 

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
+