← back to Exo
New worker
bb7f1ae99421f3e9d075c0a9ffbb6e5ce9faf8dd · 2025-07-18 10:08:56 +0100 · Alex Cheema
Co-authored-by: Matt Beton <matthew.beton@gmail.com>
Files touched
M .github/actions/typecheck/action.ymlM engines/mlx/auto_parallel.pyM engines/mlx/utils_mlx.pyM shared/types/api.pyM shared/types/events/common.pyM shared/types/events/events.pyM shared/types/states/master.pyM shared/types/states/shared.pyM shared/types/states/worker.pyM shared/types/tasks/common.pyM shared/types/worker/commands_runner.pyM shared/types/worker/downloads.pyM shared/types/worker/instances.pyM shared/types/worker/mlx.pyA shared/types/worker/ops.pyM shared/types/worker/runners.pyM shared/types/worker/shards.pyA test_shard_serialization.pyM worker/main.pyM worker/runner/communication.pyM worker/runner/runner.pyM worker/runner/runner_supervisor.pyA worker/test_worker_state.pyM worker/tests/conftest.pyM worker/tests/test_serdes.pyM worker/tests/test_supervisor.pyA worker/tests/test_worker_handlers.pyA worker/tests/test_worker_plan.py
Diff
commit bb7f1ae99421f3e9d075c0a9ffbb6e5ce9faf8dd
Author: Alex Cheema <41707476+AlexCheema@users.noreply.github.com>
Date: Fri Jul 18 10:08:56 2025 +0100
New worker
Co-authored-by: Matt Beton <matthew.beton@gmail.com>
---
.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 | 153 ++++++++++----
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, 1367 insertions(+), 235 deletions(-)
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,7 +45,55 @@ 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)
-class StartingRunnerStatus(RunnerStatus[RunnerStatusType.Starting]):
- runner_status: Literal[RunnerStatusType.Starting]
+# Emitted by the Worker
+class DownloadingRunnerStatus(BaseRunnerStatus[RunnerStatusType.Downloading]):
+ runner_status: Literal[RunnerStatusType.Downloading] = Field(default=RunnerStatusType.Downloading)
+ download_progress: DownloadProgress
+# Emitted by the Worker
+class ReadyRunnerStatus(BaseRunnerStatus[RunnerStatusType.Ready]):
+ runner_status: Literal[RunnerStatusType.Ready] = Field(default=RunnerStatusType.Ready)
-class DownloadingRunnerStatus(RunnerStatus[RunnerStatusType.Downloading]):
- runner_status: Literal[RunnerStatusType.Downloading]
- download_progress: BaseDownloadProgress[DownloadStatus]
+# Emitted by the Master
+class StartingRunnerStatus(BaseRunnerStatus[RunnerStatusType.Starting]):
+ runner_status: Literal[RunnerStatusType.Starting] = Field(default=RunnerStatusType.Starting)
+# Emitted by the Worker
+class LoadedRunnerStatus(BaseRunnerStatus[RunnerStatusType.Loaded]):
+ runner_status: Literal[RunnerStatusType.Loaded] = Field(default=RunnerStatusType.Loaded)
-class RunningRunnerStatus(RunnerStatus[RunnerStatusType.Running]):
- runner_status: Literal[RunnerStatusType.Running]
+# Emitted by the Worker
+class RunningRunnerStatus(BaseRunnerStatus[RunnerStatusType.Running]):
+ runner_status: Literal[RunnerStatusType.Running] = Field(default=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
+# ---------------------------------------------------------------------------
+
+# "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[PartitionStrategyT: PartitionStrategy](BaseModel):
+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_streaming_task_data(
- chat_completion_params: CompletionCreateParamsStreaming,
-):
- """Creates ChatCompletionStreamingTask from params"""
- return ChatCompletionStreamingTask(task_data=chat_completion_params)
-
+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 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
+
← cc45c7e9 Fixed events issue.
·
back to Exo
·
Add sqlite connector 2f64e30d →