← back to Exo
new tagged union
1c6b5ce911cf012aaf6a9df918fa2207f866ad57 · 2025-10-10 16:22:09 +0100 · Evan Quiney
Co-authored-by: Alex Cheema <alexcheema123@gmail.com>
Sorry Andrei!
Files touched
M dashboard/index.htmlM src/exo/master/api.pyM src/exo/master/main.pyM src/exo/master/placement.pyD src/exo/master/tests/api_utils_test.pyM src/exo/master/tests/test_api.pyM src/exo/master/tests/test_master.pyM src/exo/master/tests/test_placement.pyM src/exo/shared/apply.pyM src/exo/shared/types/chunks.pyM src/exo/shared/types/commands.pyM src/exo/shared/types/common.pyM src/exo/shared/types/events.pyM src/exo/shared/types/models.pyM src/exo/shared/types/state.pyM src/exo/shared/types/tasks.pyM src/exo/shared/types/worker/commands_runner.pyM src/exo/shared/types/worker/common.pyM src/exo/shared/types/worker/communication.pyM src/exo/shared/types/worker/downloads.pyM src/exo/shared/types/worker/instances.pyM src/exo/shared/types/worker/ops.pyM src/exo/shared/types/worker/runners.pyM src/exo/shared/types/worker/shards.pyM src/exo/utils/pydantic_ext.pyD src/exo/utils/pydantic_tagged.pyM src/exo/utils/tests/test_tagged.pyM src/exo/worker/download/download_utils.pyM src/exo/worker/download/impl_shard_downloader.pyM src/exo/worker/download/shard_downloader.pyM src/exo/worker/main.pyM src/exo/worker/plan.pyM src/exo/worker/tests/conftest.pyM src/exo/worker/tests/test_handlers/test_handlers_happy.pyM src/exo/worker/tests/test_integration/test_inference.pyM src/exo/worker/tests/test_integration/test_inference_sad.pyM src/exo/worker/tests/test_integration/test_instantiation.pyM src/exo/worker/tests/test_integration/test_instantiation_sad.pyM src/exo/worker/tests/test_multimodel/test_inference_llama70B.pyM src/exo/worker/tests/test_plan/test_worker_plan.pyM src/exo/worker/tests/test_plan/test_worker_plan_utils.pyM src/exo/worker/tests/test_runner_connection.pyM src/exo/worker/tests/test_serdes.pyM src/exo/worker/tests/test_supervisor/test_supervisor.pyM src/exo/worker/tests/worker_management.py
Diff
commit 1c6b5ce911cf012aaf6a9df918fa2207f866ad57
Author: Evan Quiney <evanev7@gmail.com>
Date: Fri Oct 10 16:22:09 2025 +0100
new tagged union
Co-authored-by: Alex Cheema <alexcheema123@gmail.com>
Sorry Andrei!
---
dashboard/index.html | 72 +++--
src/exo/master/api.py | 7 +-
src/exo/master/main.py | 18 +-
src/exo/master/placement.py | 2 +-
src/exo/master/tests/api_utils_test.py | 86 ------
src/exo/master/tests/test_api.py | 9 +-
src/exo/master/tests/test_master.py | 251 +++++++++---------
src/exo/master/tests/test_placement.py | 2 +-
src/exo/shared/apply.py | 6 +-
src/exo/shared/types/chunks.py | 19 +-
src/exo/shared/types/commands.py | 36 +--
src/exo/shared/types/common.py | 12 +-
src/exo/shared/types/events.py | 76 ++----
src/exo/shared/types/models.py | 4 +-
src/exo/shared/types/state.py | 18 +-
src/exo/shared/types/tasks.py | 27 +-
src/exo/shared/types/worker/commands_runner.py | 83 ++----
src/exo/shared/types/worker/common.py | 6 +-
src/exo/shared/types/worker/communication.py | 3 -
src/exo/shared/types/worker/downloads.py | 55 +---
src/exo/shared/types/worker/instances.py | 9 +-
src/exo/shared/types/worker/ops.py | 71 ++---
src/exo/shared/types/worker/runners.py | 62 ++---
src/exo/shared/types/worker/shards.py | 49 +---
src/exo/utils/pydantic_ext.py | 27 +-
src/exo/utils/pydantic_tagged.py | 229 ----------------
src/exo/utils/tests/test_tagged.py | 292 +++++++++++++--------
src/exo/worker/download/download_utils.py | 8 +-
src/exo/worker/download/impl_shard_downloader.py | 3 -
src/exo/worker/download/shard_downloader.py | 3 -
src/exo/worker/main.py | 50 ++--
src/exo/worker/plan.py | 40 ++-
src/exo/worker/tests/conftest.py | 6 +-
.../tests/test_handlers/test_handlers_happy.py | 4 +-
.../tests/test_integration/test_inference.py | 10 +-
.../tests/test_integration/test_inference_sad.py | 42 +--
.../tests/test_integration/test_instantiation.py | 8 +-
.../test_integration/test_instantiation_sad.py | 16 +-
.../test_multimodel/test_inference_llama70B.py | 25 +-
src/exo/worker/tests/test_plan/test_worker_plan.py | 54 ++--
.../tests/test_plan/test_worker_plan_utils.py | 11 +-
src/exo/worker/tests/test_runner_connection.py | 2 +-
src/exo/worker/tests/test_serdes.py | 6 +-
.../tests/test_supervisor/test_supervisor.py | 4 +-
src/exo/worker/tests/worker_management.py | 22 +-
45 files changed, 693 insertions(+), 1152 deletions(-)
diff --git a/dashboard/index.html b/dashboard/index.html
index 7560298f..23b4acb9 100644
--- a/dashboard/index.html
+++ b/dashboard/index.html
@@ -934,7 +934,7 @@
headers: {
'Content-Type': 'application/json',
},
- body: JSON.stringify({ model_id: selectedModelId })
+ body: JSON.stringify({ modelId: selectedModelId, model_id: selectedModelId })
});
if (!response.ok) {
@@ -976,26 +976,43 @@
// Calculate download status for an instance based on its runners
function calculateInstanceDownloadStatus(instance, runners) {
- if (!instance.shard_assignments?.runner_to_shard || !runners) {
+ const shardAssignments = instance.shard_assignments ?? instance.shardAssignments;
+ const runnerToShard = shardAssignments?.runner_to_shard ?? shardAssignments?.runnerToShard;
+ if (!runnerToShard || !runners) {
return { isDownloading: false, progress: 0 };
}
- const runnerIds = Object.keys(instance.shard_assignments.runner_to_shard);
+ const runnerIds = Object.keys(runnerToShard);
const downloadingRunners = [];
let totalBytes = 0;
let downloadedBytes = 0;
for (const runnerId of runnerIds) {
const runner = runners[runnerId];
+ let isRunnerDownloading = false;
+
+ // Legacy snake_case structure
if (runner && runner.runner_status === 'Downloading' && runner.download_progress) {
- downloadingRunners.push(runner);
-
- // Aggregate download progress across all downloading runners
- if (runner.download_progress.download_status === 'Downloading' && runner.download_progress.download_progress) {
+ isRunnerDownloading = runner.download_progress.download_status === 'Downloading';
+ if (isRunnerDownloading && runner.download_progress.download_progress) {
totalBytes += runner.download_progress.download_progress.total_bytes || 0;
downloadedBytes += runner.download_progress.download_progress.downloaded_bytes || 0;
}
+ } else if (runner && typeof runner === 'object') {
+ // Tagged-union camelCase structure, e.g. { "DownloadingRunnerStatus": { downloadProgress: { totalBytes, downloadedBytes } } }
+ const tag = Object.keys(runner)[0];
+ if (tag && /DownloadingRunnerStatus$/i.test(tag)) {
+ isRunnerDownloading = true;
+ const inner = runner[tag] || {};
+ const prog = inner.downloadProgress || inner.download_progress || {};
+ const t = prog.totalBytes ?? prog.total_bytes ?? 0;
+ const d = prog.downloadedBytes ?? prog.downloaded_bytes ?? 0;
+ totalBytes += typeof t === 'number' ? t : 0;
+ downloadedBytes += typeof d === 'number' ? d : 0;
+ }
}
+
+ if (isRunnerDownloading) downloadingRunners.push(runner);
}
const isDownloading = downloadingRunners.length > 0;
@@ -1007,16 +1024,25 @@
// Derive a display status for an instance from its runners.
// Priority: FAILED > DOWNLOADING > STARTING > RUNNING > LOADED > INACTIVE
function deriveInstanceStatus(instance, runners = {}) {
- const runnerIds = Object.keys(instance.shard_assignments?.runner_to_shard || {});
+ const shardAssignments = instance.shard_assignments ?? instance.shardAssignments;
+ const runnerToShard = shardAssignments?.runner_to_shard ?? shardAssignments?.runnerToShard ?? {};
+ const runnerIds = Object.keys(runnerToShard);
const statuses = runnerIds
- .map(rid => runners[rid]?.runner_status)
+ .map(rid => {
+ const r = runners[rid];
+ if (!r || typeof r !== 'object') return undefined;
+ if (typeof r.runner_status === 'string') return r.runner_status;
+ const tag = Object.keys(r)[0];
+ return typeof tag === 'string' ? tag.replace(/RunnerStatus$/,'') : undefined; // e.g. LoadedRunnerStatus -> Loaded
+ })
.filter(s => typeof s === 'string');
const has = (s) => statuses.includes(s);
const every = (pred) => statuses.length > 0 && statuses.every(pred);
if (statuses.length === 0) {
- const inactive = instance.instance_type === 'INACTIVE';
+ const instanceType = instance.instance_type ?? instance.instanceType;
+ const inactive = instanceType === 'INACTIVE' || instanceType === 'Inactive';
return { statusText: inactive ? 'INACTIVE' : 'LOADED', statusClass: inactive ? 'inactive' : 'loaded' };
}
@@ -1046,10 +1072,12 @@
}
const instancesHTML = instancesArray.map(instance => {
- const modelId = instance.shard_assignments?.model_id || 'Unknown Model';
- const truncatedInstanceId = instance.instance_id.length > 8
- ? instance.instance_id.substring(0, 8) + '...'
- : instance.instance_id;
+ const shardAssignments = instance.shard_assignments ?? instance.shardAssignments;
+ const modelId = shardAssignments?.model_id ?? shardAssignments?.modelId ?? 'Unknown Model';
+ const instanceId = instance.instance_id ?? instance.instanceId ?? '';
+ const truncatedInstanceId = instanceId.length > 8
+ ? instanceId.substring(0, 8) + '...'
+ : instanceId;
const hostsHTML = instance.hosts?.map(host =>
`<span class="instance-host">${host.ip}:${host.port}</span>`
@@ -1083,14 +1111,14 @@
<span class="instance-status ${statusClass}">${statusText}</span>
</div>
<div class="instance-actions">
- <button class="instance-delete-button" data-instance-id="${instance.instance_id}" title="Delete Instance">
+ <button class="instance-delete-button" data-instance-id="${instanceId}" title="Delete Instance">
Delete
</button>
</div>
</div>
<div class="instance-model">${modelId}</div>
<div class="instance-details">
- Shards: ${Object.keys(instance.shard_assignments?.runner_to_shard || {}).length}
+ Shards: ${Object.keys((shardAssignments?.runner_to_shard ?? shardAssignments?.runnerToShard) || {}).length}
</div>
${downloadProgressHTML}
${hostsHTML ? `<div class="instance-hosts">${hostsHTML}</div>` : ''}
@@ -1931,16 +1959,18 @@
return fallback;
};
- // Process nodes from topology or fallback to node_profiles directly
+ // Process nodes from topology or fallback to node_profiles/nodeProfiles directly
let nodesToProcess = {};
if (clusterState.topology && Array.isArray(clusterState.topology.nodes)) {
clusterState.topology.nodes.forEach(node => {
- if (node.node_id && node.node_profile) {
- nodesToProcess[node.node_id] = node.node_profile;
+ const nid = node.node_id ?? node.nodeId;
+ const nprof = node.node_profile ?? node.nodeProfile;
+ if (nid && nprof) {
+ nodesToProcess[nid] = nprof;
}
});
- } else if (clusterState.node_profiles) {
- nodesToProcess = clusterState.node_profiles;
+ } else if (clusterState.node_profiles || clusterState.nodeProfiles) {
+ nodesToProcess = clusterState.node_profiles ?? clusterState.nodeProfiles;
}
// Transform each node
diff --git a/src/exo/master/api.py b/src/exo/master/api.py
index d10f7dd6..a4ad65cd 100644
--- a/src/exo/master/api.py
+++ b/src/exo/master/api.py
@@ -33,7 +33,6 @@ from exo.shared.types.commands import (
CreateInstance,
DeleteInstance,
ForwarderCommand,
- TaggedCommand,
# TODO: SpinUpInstance
TaskFinished,
)
@@ -306,7 +305,7 @@ class API:
async def _apply_state(self):
with self.global_event_receiver as events:
async for event in events:
- self.event_buffer.ingest(event.origin_idx, event.tagged_event.c)
+ self.event_buffer.ingest(event.origin_idx, event.event)
for idx, event in self.event_buffer.drain_indexed():
self.state = apply(self.state, IndexedEvent(event=event, idx=idx))
if (
@@ -317,7 +316,5 @@ class API:
async def _send(self, command: Command):
await self.command_sender.send(
- ForwarderCommand(
- origin=self.node_id, tagged_command=TaggedCommand.from_(command)
- )
+ ForwarderCommand(origin=self.node_id, command=command)
)
diff --git a/src/exo/master/main.py b/src/exo/master/main.py
index 443a2803..ce3643c2 100644
--- a/src/exo/master/main.py
+++ b/src/exo/master/main.py
@@ -23,13 +23,12 @@ from exo.shared.types.events import (
ForwarderEvent,
IndexedEvent,
InstanceDeleted,
- TaggedEvent,
TaskCreated,
TaskDeleted,
TopologyEdgeDeleted,
)
from exo.shared.types.state import State
-from exo.shared.types.tasks import ChatCompletionTask, TaskId, TaskStatus, TaskType
+from exo.shared.types.tasks import ChatCompletionTask, TaskId, TaskStatus
from exo.shared.types.worker.common import InstanceId
from exo.utils.channels import Receiver, Sender, channel
from exo.utils.event_buffer import MultiSourceBuffer
@@ -90,11 +89,9 @@ class Master:
with self.command_receiver as commands:
async for forwarder_command in commands:
try:
- logger.info(
- f"Executing command: {forwarder_command.tagged_command.c}"
- )
+ logger.info(f"Executing command: {forwarder_command.command}")
generated_events: list[Event] = []
- command = forwarder_command.tagged_command.c
+ command = forwarder_command.command
match command:
case ChatCompletion():
instance_task_counts: dict[InstanceId, int] = {}
@@ -130,11 +127,10 @@ class Master:
TaskCreated(
task_id=task_id,
task=ChatCompletionTask(
- task_type=TaskType.CHAT_COMPLETION,
task_id=task_id,
command_id=command.command_id,
instance_id=available_instance_ids[0],
- task_status=TaskStatus.PENDING,
+ task_status=TaskStatus.Pending,
task_params=command.request_params,
),
)
@@ -190,7 +186,7 @@ class Master:
async for local_event in local_events:
self._multi_buffer.ingest(
local_event.origin_idx,
- local_event.tagged_event.c,
+ local_event.event,
local_event.origin,
)
for event in self._multi_buffer.drain():
@@ -224,7 +220,7 @@ class Master:
ForwarderEvent(
origin=NodeId(f"master_{self.node_id}"),
origin_idx=local_index,
- tagged_event=TaggedEvent.from_(event),
+ event=event,
)
)
local_index += 1
@@ -235,6 +231,6 @@ class Master:
ForwarderEvent(
origin=self.node_id,
origin_idx=event.idx,
- tagged_event=TaggedEvent.from_(event.event),
+ event=event.event,
)
)
diff --git a/src/exo/master/placement.py b/src/exo/master/placement.py
index e3884d53..b5e402d9 100644
--- a/src/exo/master/placement.py
+++ b/src/exo/master/placement.py
@@ -83,7 +83,7 @@ def get_instance_placements_after_create(
target_instances = dict(deepcopy(current_instances))
target_instances[instance_id] = Instance(
instance_id=instance_id,
- instance_type=InstanceStatus.ACTIVE,
+ instance_type=InstanceStatus.Active,
shard_assignments=shard_assignments,
hosts=[
Host(
diff --git a/src/exo/master/tests/api_utils_test.py b/src/exo/master/tests/api_utils_test.py
deleted file mode 100644
index 3ed52c7a..00000000
--- a/src/exo/master/tests/api_utils_test.py
+++ /dev/null
@@ -1,86 +0,0 @@
-import asyncio
-import functools
-from typing import (
- Any,
- AsyncGenerator,
- Awaitable,
- Callable,
- Coroutine,
- ParamSpec,
- TypeVar,
- final,
-)
-
-import openai
-import pytest
-from openai._streaming import AsyncStream
-from openai.types.chat import (
- ChatCompletionMessageParam,
-)
-from openai.types.chat.chat_completion_chunk import ChatCompletionChunk, Choice
-
-from exo.main import main
-
-_P = ParamSpec("_P")
-_R = TypeVar("_R")
-
-OPENAI_API_KEY: str = "<YOUR_OPENAI_API_KEY>"
-OPENAI_API_URL: str = "http://0.0.0.0:8000/v1"
-
-
-def with_master_main(
- func: Callable[_P, Awaitable[_R]],
-) -> Callable[_P, Coroutine[Any, Any, _R]]:
- @pytest.mark.asyncio
- @functools.wraps(func)
- async def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
- loop = asyncio.get_running_loop()
- master_task = loop.run_in_executor(None, main)
- try:
- return await func(*args, **kwargs)
- finally:
- master_task.cancel()
- with pytest.raises(asyncio.CancelledError):
- await master_task
-
- return wrapper
-
-
-@final
-class ChatMessage:
- """Strictly-typed chat message for OpenAI API."""
-
- def __init__(self, role: str, content: str) -> None:
- self.role = role
- self.content = content
-
- def to_openai(self) -> ChatCompletionMessageParam:
- if self.role == "user":
- return {"role": "user", "content": self.content} # type: ChatCompletionUserMessageParam
- elif self.role == "assistant":
- return {"role": "assistant", "content": self.content} # type: ChatCompletionAssistantMessageParam
- elif self.role == "system":
- return {"role": "system", "content": self.content} # type: ChatCompletionSystemMessageParam
- else:
- raise ValueError(f"Unsupported role: {self.role}")
-
-
-async def stream_chatgpt_response(
- messages: list[ChatMessage],
- model: str = "gpt-3.5-turbo",
-) -> AsyncGenerator[Choice, None]:
- client = openai.AsyncOpenAI(
- api_key=OPENAI_API_KEY,
- base_url=OPENAI_API_URL,
- )
- openai_messages: list[ChatCompletionMessageParam] = [
- m.to_openai() for m in messages
- ]
- stream: AsyncStream[ChatCompletionChunk] = await client.chat.completions.create(
- model=model,
- messages=openai_messages,
- stream=True,
- )
- async for chunk in stream:
- for choice in chunk.choices:
- yield choice
diff --git a/src/exo/master/tests/test_api.py b/src/exo/master/tests/test_api.py
index ce9e1376..5965ab5e 100644
--- a/src/exo/master/tests/test_api.py
+++ b/src/exo/master/tests/test_api.py
@@ -2,17 +2,10 @@ import asyncio
import pytest
-from exo.master.tests.api_utils_test import (
- ChatMessage,
- stream_chatgpt_response,
- with_master_main,
-)
-
-@with_master_main
@pytest.mark.asyncio
async def test_master_api_multiple_response_sequential() -> None:
- # TODO: This hangs at the moment it seems.
+ # TODO
return
messages = [ChatMessage(role="user", content="Hello, who are you?")]
token_count = 0
diff --git a/src/exo/master/tests/test_master.py b/src/exo/master/tests/test_master.py
index bfa3f564..a1b6c0b6 100644
--- a/src/exo/master/tests/test_master.py
+++ b/src/exo/master/tests/test_master.py
@@ -1,7 +1,8 @@
-import asyncio
from typing import List, Sequence
+import anyio
import pytest
+from loguru import logger
from exo.master.main import Master
from exo.routing.router import get_node_id_keypair
@@ -11,7 +12,6 @@ from exo.shared.types.commands import (
CommandId,
CreateInstance,
ForwarderCommand,
- TaggedCommand,
)
from exo.shared.types.common import NodeId
from exo.shared.types.events import (
@@ -19,7 +19,6 @@ from exo.shared.types.events import (
IndexedEvent,
InstanceCreated,
NodePerformanceMeasured,
- TaggedEvent,
TaskCreated,
)
from exo.shared.types.memory import Memory
@@ -29,9 +28,9 @@ from exo.shared.types.profiling import (
NodePerformanceProfile,
SystemPerformanceProfile,
)
-from exo.shared.types.tasks import ChatCompletionTask, TaskStatus, TaskType
+from exo.shared.types.tasks import ChatCompletionTask, TaskStatus
from exo.shared.types.worker.instances import Instance, InstanceStatus, ShardAssignments
-from exo.shared.types.worker.shards import PartitionStrategy, PipelineShardMetadata
+from exo.shared.types.worker.shards import PipelineShardMetadata
from exo.utils.channels import channel
@@ -46,12 +45,12 @@ async def test_master():
all_events: List[IndexedEvent] = []
- async def _get_events() -> Sequence[IndexedEvent]:
+ def _get_events() -> Sequence[IndexedEvent]:
orig_events = global_event_receiver.collect()
for e in orig_events:
all_events.append(
IndexedEvent(
- event=e.tagged_event.c,
+ event=e.event,
idx=len(all_events), # origin=e.origin,
)
)
@@ -64,133 +63,141 @@ async def test_master():
command_receiver=co_receiver,
tb_only=False,
)
- asyncio.create_task(master.run())
+ logger.info("run the master")
+ async with anyio.create_task_group() as tg:
+ tg.start_soon(master.run)
- sender_node_id = NodeId(f"{keypair.to_peer_id().to_base58()}_sender")
- # inject a NodePerformanceProfile event
- await local_event_sender.send(
- ForwarderEvent(
- origin_idx=0,
- origin=sender_node_id,
- tagged_event=TaggedEvent.from_(
- NodePerformanceMeasured(
- node_id=node_id,
- node_profile=NodePerformanceProfile(
- model_id="maccy",
- chip_id="arm",
- friendly_name="test",
- memory=MemoryPerformanceProfile(
- ram_total=Memory.from_bytes(678948 * 1024),
- ram_available=Memory.from_bytes(678948 * 1024),
- swap_total=Memory.from_bytes(0),
- swap_available=Memory.from_bytes(0),
+ sender_node_id = NodeId(f"{keypair.to_peer_id().to_base58()}_sender")
+ # inject a NodePerformanceProfile event
+ logger.info("inject a NodePerformanceProfile event")
+ await local_event_sender.send(
+ ForwarderEvent(
+ origin_idx=0,
+ origin=sender_node_id,
+ event=(
+ NodePerformanceMeasured(
+ node_id=node_id,
+ node_profile=NodePerformanceProfile(
+ model_id="maccy",
+ chip_id="arm",
+ friendly_name="test",
+ memory=MemoryPerformanceProfile(
+ ram_total=Memory.from_bytes(678948 * 1024),
+ ram_available=Memory.from_bytes(678948 * 1024),
+ swap_total=Memory.from_bytes(0),
+ swap_available=Memory.from_bytes(0),
+ ),
+ network_interfaces=[],
+ system=SystemPerformanceProfile(flops_fp16=0),
),
- network_interfaces=[],
- system=SystemPerformanceProfile(flops_fp16=0),
- ),
- )
- ),
+ )
+ ),
+ )
)
- )
-
- # wait for initial topology event
- while len(list(master.state.topology.list_nodes())) == 0:
- await asyncio.sleep(0.001)
- while len(master.state.node_profiles) == 0:
- await asyncio.sleep(0.001)
- await command_sender.send(
- ForwarderCommand(
- origin=node_id,
- tagged_command=TaggedCommand.from_(
- CreateInstance(
- command_id=CommandId(),
- model_meta=ModelMetadata(
- model_id=ModelId("llama-3.2-1b"),
- pretty_name="Llama 3.2 1B",
- n_layers=16,
- storage_size=Memory.from_bytes(678948),
- ),
- )
- ),
- )
- )
- while len(master.state.instances.keys()) == 0:
- await asyncio.sleep(0.001)
- await command_sender.send(
- ForwarderCommand(
- origin=node_id,
- tagged_command=TaggedCommand.from_(
- ChatCompletion(
- command_id=CommandId(),
- request_params=ChatCompletionTaskParams(
- model="llama-3.2-1b",
- messages=[
- ChatCompletionMessage(
- role="user", content="Hello, how are you?"
- )
- ],
- ),
- )
- ),
- )
- )
- while len(await _get_events()) < 3:
- await asyncio.sleep(0.001)
+ # wait for initial topology event
+ logger.info("wait for initial topology event")
+ while len(list(master.state.topology.list_nodes())) == 0:
+ await anyio.sleep(0.001)
+ while len(master.state.node_profiles) == 0:
+ await anyio.sleep(0.001)
- events = await _get_events()
- assert len(events) == 3
- assert events[0].idx == 0
- assert events[1].idx == 1
- assert events[2].idx == 2
- assert isinstance(events[0].event, NodePerformanceMeasured)
- assert isinstance(events[1].event, InstanceCreated)
- runner_id = list(events[1].event.instance.shard_assignments.runner_to_shard.keys())[
- 0
- ]
- assert events[1].event == InstanceCreated(
- event_id=events[1].event.event_id,
- instance=Instance(
- instance_id=events[1].event.instance.instance_id,
- instance_type=InstanceStatus.ACTIVE,
- shard_assignments=ShardAssignments(
- model_id=ModelId("llama-3.2-1b"),
- runner_to_shard={
- (runner_id): PipelineShardMetadata(
- partition_strategy=PartitionStrategy.pipeline,
- start_layer=0,
- end_layer=16,
- n_layers=16,
+ logger.info("inject a CreateInstance Command")
+ await command_sender.send(
+ ForwarderCommand(
+ origin=node_id,
+ command=(
+ CreateInstance(
+ command_id=CommandId(),
model_meta=ModelMetadata(
model_id=ModelId("llama-3.2-1b"),
pretty_name="Llama 3.2 1B",
n_layers=16,
storage_size=Memory.from_bytes(678948),
),
- device_rank=0,
- world_size=1,
)
- },
- node_to_runner={node_id: runner_id},
+ ),
+ )
+ )
+ logger.info("wait for an instance")
+ while len(master.state.instances.keys()) == 0:
+ await anyio.sleep(0.001)
+ logger.info("inject a ChatCompletion Command")
+ await command_sender.send(
+ ForwarderCommand(
+ origin=node_id,
+ command=(
+ ChatCompletion(
+ command_id=CommandId(),
+ request_params=ChatCompletionTaskParams(
+ model="llama-3.2-1b",
+ messages=[
+ ChatCompletionMessage(
+ role="user", content="Hello, how are you?"
+ )
+ ],
+ ),
+ )
+ ),
+ )
+ )
+ while len(_get_events()) < 3:
+ await anyio.sleep(0.01)
+
+ events = _get_events()
+ assert len(events) == 3
+ assert events[0].idx == 0
+ assert events[1].idx == 1
+ assert events[2].idx == 2
+ assert isinstance(events[0].event, NodePerformanceMeasured)
+ assert isinstance(events[1].event, InstanceCreated)
+ runner_id = list(
+ events[1].event.instance.shard_assignments.runner_to_shard.keys()
+ )[0]
+ assert events[1].event == InstanceCreated(
+ event_id=events[1].event.event_id,
+ instance=Instance(
+ instance_id=events[1].event.instance.instance_id,
+ instance_type=InstanceStatus.Active,
+ shard_assignments=ShardAssignments(
+ model_id=ModelId("llama-3.2-1b"),
+ runner_to_shard={
+ (runner_id): PipelineShardMetadata(
+ start_layer=0,
+ end_layer=16,
+ n_layers=16,
+ model_meta=ModelMetadata(
+ model_id=ModelId("llama-3.2-1b"),
+ pretty_name="Llama 3.2 1B",
+ n_layers=16,
+ storage_size=Memory.from_bytes(678948),
+ ),
+ device_rank=0,
+ world_size=1,
+ )
+ },
+ node_to_runner={node_id: runner_id},
+ ),
+ hosts=[],
),
- hosts=[],
- ),
- )
- assert isinstance(events[2].event, TaskCreated)
- assert events[2].event == TaskCreated(
- event_id=events[2].event.event_id,
- task_id=events[2].event.task_id,
- task=ChatCompletionTask(
+ )
+ assert isinstance(events[2].event, TaskCreated)
+ assert events[2].event == TaskCreated(
+ event_id=events[2].event.event_id,
task_id=events[2].event.task_id,
- command_id=events[2].event.task.command_id,
- task_type=TaskType.CHAT_COMPLETION,
- instance_id=events[2].event.task.instance_id,
- task_status=TaskStatus.PENDING,
- task_params=ChatCompletionTaskParams(
- model="llama-3.2-1b",
- messages=[
- ChatCompletionMessage(role="user", content="Hello, how are you?")
- ],
+ task=ChatCompletionTask(
+ task_id=events[2].event.task_id,
+ command_id=events[2].event.task.command_id,
+ instance_id=events[2].event.task.instance_id,
+ task_status=TaskStatus.Pending,
+ task_params=ChatCompletionTaskParams(
+ model="llama-3.2-1b",
+ messages=[
+ ChatCompletionMessage(
+ role="user", content="Hello, how are you?"
+ )
+ ],
+ ),
),
- ),
- )
+ )
+ await master.shutdown()
diff --git a/src/exo/master/tests/test_placement.py b/src/exo/master/tests/test_placement.py
index 6b3aabf6..d210c9ff 100644
--- a/src/exo/master/tests/test_placement.py
+++ b/src/exo/master/tests/test_placement.py
@@ -27,7 +27,7 @@ def topology() -> Topology:
def instance() -> Instance:
return Instance(
instance_id=InstanceId(),
- instance_type=InstanceStatus.ACTIVE,
+ instance_type=InstanceStatus.Active,
shard_assignments=ShardAssignments(
model_id=ModelId("test-model"), runner_to_shard={}, node_to_runner={}
),
diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py
index 1ba73d7b..08150783 100644
--- a/src/exo/shared/apply.py
+++ b/src/exo/shared/apply.py
@@ -104,7 +104,7 @@ def apply_task_state_updated(event: TaskStateUpdated, state: State) -> State:
update: dict[str, TaskStatus | None] = {
"task_status": event.task_status,
}
- if event.task_status != TaskStatus.FAILED:
+ if event.task_status != TaskStatus.Failed:
update["error_type"] = None
update["error_message"] = None
@@ -138,7 +138,7 @@ def apply_instance_activated(event: InstanceActivated, state: State) -> State:
return state
updated_instance = state.instances[event.instance_id].model_copy(
- update={"instance_type": InstanceStatus.ACTIVE}
+ update={"instance_type": InstanceStatus.Active}
)
new_instances: Mapping[InstanceId, Instance] = {
**state.instances,
@@ -152,7 +152,7 @@ def apply_instance_deactivated(event: InstanceDeactivated, state: State) -> Stat
return state
updated_instance = state.instances[event.instance_id].model_copy(
- update={"instance_type": InstanceStatus.INACTIVE}
+ update={"instance_type": InstanceStatus.Inactive}
)
new_instances: Mapping[InstanceId, Instance] = {
**state.instances,
diff --git a/src/exo/shared/types/chunks.py b/src/exo/shared/types/chunks.py
index ec7a8295..f74c901a 100644
--- a/src/exo/shared/types/chunks.py
+++ b/src/exo/shared/types/chunks.py
@@ -1,35 +1,30 @@
from enum import Enum
-from typing import Annotated, Literal
-
-from pydantic import BaseModel, Field
from exo.shared.openai_compat import FinishReason
from exo.shared.types.common import CommandId
from exo.shared.types.models import ModelId
+from exo.utils.pydantic_ext import TaggedModel
class ChunkType(str, Enum):
- token = "token"
- image = "image"
+ Token = "Token"
+ Image = "Image"
-class BaseChunk[ChunkTypeT: ChunkType](BaseModel):
- chunk_type: ChunkTypeT
+class BaseChunk(TaggedModel):
command_id: CommandId
idx: int
model: ModelId
-class TokenChunk(BaseChunk[ChunkType.token]):
- chunk_type: Literal[ChunkType.token] = Field(default=ChunkType.token, frozen=True)
+class TokenChunk(BaseChunk):
text: str
token_id: int
finish_reason: FinishReason | None = None
-class ImageChunk(BaseChunk[ChunkType.image]):
- chunk_type: Literal[ChunkType.image] = Field(default=ChunkType.image, frozen=True)
+class ImageChunk(BaseChunk):
data: bytes
-GenerationChunk = Annotated[TokenChunk | ImageChunk, Field(discriminator="chunk_type")]
+GenerationChunk = TokenChunk | ImageChunk
diff --git a/src/exo/shared/types/commands.py b/src/exo/shared/types/commands.py
index 4c9a1066..d7f5da87 100644
--- a/src/exo/shared/types/commands.py
+++ b/src/exo/shared/types/commands.py
@@ -1,5 +1,4 @@
from enum import Enum
-from typing import Union
from pydantic import Field
@@ -7,8 +6,7 @@ from exo.shared.types.api import ChatCompletionTaskParams
from exo.shared.types.common import CommandId, NodeId
from exo.shared.types.models import ModelMetadata
from exo.shared.types.worker.common import InstanceId
-from exo.utils.pydantic_ext import CamelCaseModel
-from exo.utils.pydantic_tagged import Tagged, tagged_union
+from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
# TODO: We need to have a distinction between create instance and spin up instance.
@@ -21,7 +19,7 @@ class CommandType(str, Enum):
RequestEventLog = "RequestEventLog"
-class BaseCommand(CamelCaseModel):
+class BaseCommand(TaggedModel):
command_id: CommandId = Field(default_factory=CommandId)
@@ -49,30 +47,16 @@ class RequestEventLog(BaseCommand):
since_idx: int
-Command = Union[
- RequestEventLog,
- ChatCompletion,
- CreateInstance,
- SpinUpInstance,
- DeleteInstance,
- TaskFinished,
-]
-
-
-@tagged_union(
- {
- CommandType.ChatCompletion: ChatCompletion,
- CommandType.CreateInstance: CreateInstance,
- CommandType.SpinUpInstance: SpinUpInstance,
- CommandType.DeleteInstance: DeleteInstance,
- CommandType.TaskFinished: TaskFinished,
- CommandType.RequestEventLog: RequestEventLog,
- }
+Command = (
+ RequestEventLog
+ | ChatCompletion
+ | CreateInstance
+ | SpinUpInstance
+ | DeleteInstance
+ | TaskFinished
)
-class TaggedCommand(Tagged[Command]):
- pass
class ForwarderCommand(CamelCaseModel):
origin: NodeId
- tagged_command: TaggedCommand
+ command: Command
diff --git a/src/exo/shared/types/common.py b/src/exo/shared/types/common.py
index b89ff915..e34fc7ef 100644
--- a/src/exo/shared/types/common.py
+++ b/src/exo/shared/types/common.py
@@ -1,11 +1,13 @@
from typing import Self
from uuid import uuid4
-from pydantic import BaseModel, GetCoreSchemaHandler, field_validator
+from pydantic import GetCoreSchemaHandler, field_validator
from pydantic_core import core_schema
+from exo.utils.pydantic_ext import CamelCaseModel
-class ID(str):
+
+class Id(str):
def __new__(cls, value: str | None = None) -> Self:
return super().__new__(cls, value or str(uuid4()))
@@ -17,15 +19,15 @@ class ID(str):
return core_schema.str_schema()
-class NodeId(ID):
+class NodeId(Id):
pass
-class CommandId(ID):
+class CommandId(Id):
pass
-class Host(BaseModel):
+class Host(CamelCaseModel):
ip: str
port: int
diff --git a/src/exo/shared/types/events.py b/src/exo/shared/types/events.py
index 074457a3..a910ea93 100644
--- a/src/exo/shared/types/events.py
+++ b/src/exo/shared/types/events.py
@@ -1,21 +1,19 @@
from enum import Enum
-from typing import Union
from pydantic import Field
from exo.shared.topology import Connection, NodePerformanceProfile
from exo.shared.types.chunks import CommandId, GenerationChunk
-from exo.shared.types.common import ID, NodeId
+from exo.shared.types.common import Id, NodeId
from exo.shared.types.profiling import MemoryPerformanceProfile
from exo.shared.types.tasks import Task, TaskId, TaskStatus
from exo.shared.types.worker.common import InstanceId, WorkerStatus
from exo.shared.types.worker.instances import Instance
from exo.shared.types.worker.runners import RunnerId, RunnerStatus
-from exo.utils.pydantic_ext import CamelCaseModel
-from exo.utils.pydantic_tagged import Tagged, tagged_union
+from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
-class EventId(ID):
+class EventId(Id):
"""
Newtype around `ID`
"""
@@ -60,7 +58,7 @@ class EventType(str, Enum):
TopologyEdgeDeleted = "TopologyEdgeDeleted"
-class BaseEvent(CamelCaseModel):
+class BaseEvent(TaggedModel):
event_id: EventId = Field(default_factory=EventId)
@@ -145,52 +143,26 @@ class TopologyEdgeDeleted(BaseEvent):
edge: Connection
-Event = Union[
- TestEvent,
- TaskCreated,
- TaskStateUpdated,
- TaskFailed,
- TaskDeleted,
- InstanceCreated,
- InstanceActivated,
- InstanceDeactivated,
- InstanceDeleted,
- RunnerStatusUpdated,
- RunnerDeleted,
- NodePerformanceMeasured,
- NodeMemoryMeasured,
- WorkerStatusUpdated,
- ChunkGenerated,
- TopologyNodeCreated,
- TopologyEdgeCreated,
- TopologyEdgeDeleted,
-]
-
-
-@tagged_union(
- {
- EventType.TestEvent: TestEvent,
- EventType.TaskCreated: TaskCreated,
- EventType.TaskStateUpdated: TaskStateUpdated,
- EventType.TaskFailed: TaskFailed,
- EventType.TaskDeleted: TaskDeleted,
- EventType.InstanceCreated: InstanceCreated,
- EventType.InstanceActivated: InstanceActivated,
- EventType.InstanceDeactivated: InstanceDeactivated,
- EventType.InstanceDeleted: InstanceDeleted,
- EventType.RunnerStatusUpdated: RunnerStatusUpdated,
- EventType.RunnerDeleted: RunnerDeleted,
- EventType.NodePerformanceMeasured: NodePerformanceMeasured,
- EventType.NodeMemoryMeasured: NodeMemoryMeasured,
- EventType.WorkerStatusUpdated: WorkerStatusUpdated,
- EventType.ChunkGenerated: ChunkGenerated,
- EventType.TopologyNodeCreated: TopologyNodeCreated,
- EventType.TopologyEdgeCreated: TopologyEdgeCreated,
- EventType.TopologyEdgeDeleted: TopologyEdgeDeleted,
- }
+Event = (
+ TestEvent
+ | TaskCreated
+ | TaskStateUpdated
+ | TaskFailed
+ | TaskDeleted
+ | InstanceCreated
+ | InstanceActivated
+ | InstanceDeactivated
+ | InstanceDeleted
+ | RunnerStatusUpdated
+ | RunnerDeleted
+ | NodePerformanceMeasured
+ | NodeMemoryMeasured
+ | WorkerStatusUpdated
+ | ChunkGenerated
+ | TopologyNodeCreated
+ | TopologyEdgeCreated
+ | TopologyEdgeDeleted
)
-class TaggedEvent(Tagged[Event]):
- pass
class IndexedEvent(CamelCaseModel):
@@ -205,4 +177,4 @@ class ForwarderEvent(CamelCaseModel):
origin_idx: int = Field(ge=0)
origin: NodeId
- tagged_event: TaggedEvent
+ event: Event
diff --git a/src/exo/shared/types/models.py b/src/exo/shared/types/models.py
index eaff0d79..b029fba0 100644
--- a/src/exo/shared/types/models.py
+++ b/src/exo/shared/types/models.py
@@ -1,11 +1,11 @@
from pydantic import PositiveInt
-from exo.shared.types.common import ID
+from exo.shared.types.common import Id
from exo.shared.types.memory import Memory
from exo.utils.pydantic_ext import CamelCaseModel
-class ModelId(ID):
+class ModelId(Id):
pass
diff --git a/src/exo/shared/types/state.py b/src/exo/shared/types/state.py
index e599b0af..8e2e6ede 100644
--- a/src/exo/shared/types/state.py
+++ b/src/exo/shared/types/state.py
@@ -1,7 +1,7 @@
from collections.abc import Mapping, Sequence
from typing import Any, cast
-from pydantic import BaseModel, ConfigDict, Field, field_validator
+from pydantic import ConfigDict, Field, field_validator, field_serializer
from exo.shared.topology import Topology, TopologySnapshot
from exo.shared.types.common import NodeId
@@ -10,15 +10,10 @@ from exo.shared.types.tasks import Task, TaskId
from exo.shared.types.worker.common import InstanceId, WorkerStatus
from exo.shared.types.worker.instances import Instance
from exo.shared.types.worker.runners import RunnerId, RunnerStatus
+from exo.utils.pydantic_ext import CamelCaseModel
-def _encode_topology(topo: "Topology") -> dict[str, Any]: # noqa: D401
- """Serialise *topo* into a JSON-compatible dict."""
-
- return topo.to_snapshot().model_dump()
-
-
-class State(BaseModel):
+class State(CamelCaseModel):
"""Global system state.
The :class:`Topology` instance is encoded/decoded via an immutable
@@ -28,9 +23,6 @@ class State(BaseModel):
model_config = ConfigDict(
arbitrary_types_allowed=True,
- json_encoders={
- Topology: _encode_topology,
- },
)
node_status: Mapping[NodeId, WorkerStatus] = {}
instances: Mapping[InstanceId, Instance] = {}
@@ -41,6 +33,10 @@ class State(BaseModel):
history: Sequence[Topology] = []
last_event_applied_idx: int = Field(default=-1, ge=-1)
+ @field_serializer("topology", mode="plain")
+ def _encode_topology(self, value: Topology) -> TopologySnapshot:
+ return value.to_snapshot()
+
@field_validator("topology", mode="before")
@classmethod
def _deserialize_topology(cls, value: object) -> Topology: # noqa: D401 – Pydantic validator signature
diff --git a/src/exo/shared/types/tasks.py b/src/exo/shared/types/tasks.py
index 200cef1c..c500a569 100644
--- a/src/exo/shared/types/tasks.py
+++ b/src/exo/shared/types/tasks.py
@@ -1,30 +1,25 @@
from enum import Enum
-from typing import Annotated, Literal
-from pydantic import BaseModel, Field
+from pydantic import Field
from exo.shared.types.api import ChatCompletionTaskParams
-from exo.shared.types.common import ID, CommandId
+from exo.shared.types.common import CommandId, Id
from exo.shared.types.worker.common import InstanceId
+from exo.utils.pydantic_ext import TaggedModel
-class TaskId(ID):
+class TaskId(Id):
pass
-
-
-class TaskType(str, Enum):
- CHAT_COMPLETION = "CHAT_COMPLETION"
-
+
class TaskStatus(str, Enum):
- PENDING = "PENDING"
- RUNNING = "RUNNING"
- COMPLETE = "COMPLETE"
- FAILED = "FAILED"
+ Pending = "Pending"
+ Running = "Running"
+ Complete = "Complete"
+ Failed = "Failed"
-class ChatCompletionTask(BaseModel):
- task_type: Literal[TaskType.CHAT_COMPLETION] = TaskType.CHAT_COMPLETION
+class ChatCompletionTask(TaggedModel):
task_id: TaskId
command_id: CommandId
instance_id: InstanceId
@@ -35,4 +30,4 @@ class ChatCompletionTask(BaseModel):
error_message: str | None = Field(default=None)
-Task = Annotated[ChatCompletionTask, Field(discriminator="task_type")]
+Task = ChatCompletionTask
diff --git a/src/exo/shared/types/worker/commands_runner.py b/src/exo/shared/types/worker/commands_runner.py
index 66696482..407ea2f4 100644
--- a/src/exo/shared/types/worker/commands_runner.py
+++ b/src/exo/shared/types/worker/commands_runner.py
@@ -1,116 +1,69 @@
-from enum import Enum
-from typing import Annotated, Literal
-
-from pydantic import BaseModel, Field, TypeAdapter
-
from exo.shared.openai_compat import FinishReason
from exo.shared.types.common import Host
from exo.shared.types.tasks import ChatCompletionTaskParams
from exo.shared.types.worker.shards import ShardMetadata
+from exo.utils.pydantic_ext import TaggedModel
-## Messages passed TO the runner
-class MessageType(str, Enum):
- Setup = "setup"
- ChatTask = "chat_task"
- Exit = "exit"
-
-
-class BaseRunnerMessage[MT: MessageType](BaseModel):
+class BaseRunnerMessage(TaggedModel):
pass
-class SetupMessage(BaseRunnerMessage[MessageType.Setup]):
- type: Literal[MessageType.Setup] = Field(default=MessageType.Setup, frozen=True)
+class SetupMessage(BaseRunnerMessage):
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
- )
+class ChatTaskMessage(BaseRunnerMessage):
task_data: ChatCompletionTaskParams
-class ExitMessage(BaseRunnerMessage[MessageType.Exit]):
- type: Literal[MessageType.Exit] = Field(default=MessageType.Exit, frozen=True)
-
-
-RunnerMessage = Annotated[
- SetupMessage | ChatTaskMessage | ExitMessage, Field(discriminator="type")
-]
-RunnerMessageTypeAdapter: TypeAdapter[RunnerMessage] = TypeAdapter(RunnerMessage)
+class ExitMessage(BaseRunnerMessage):
+ pass
-## Responses passed FROM the runner
-class RunnerResponseType(str, Enum):
- InitializedResponse = "initialized_response"
- TokenizedResponse = "tokenized_response"
- GenerationResponse = "generation_response"
- FinishedResponse = "finished_response"
- PrintResponse = "print_response"
- ErrorResponse = "error_response"
+RunnerMessage = SetupMessage | ChatTaskMessage | ExitMessage
-class BaseRunnerResponse[RRT: RunnerResponseType](BaseModel):
+class BaseRunnerResponse(TaggedModel):
pass
-class InitializedResponse(BaseRunnerResponse[RunnerResponseType.InitializedResponse]):
- type: Literal[RunnerResponseType.InitializedResponse] = Field(
- default=RunnerResponseType.InitializedResponse, frozen=True
- )
+class InitializedResponse(BaseRunnerResponse):
time_taken: float
-class TokenizedResponse(BaseRunnerResponse[RunnerResponseType.TokenizedResponse]):
- type: Literal[RunnerResponseType.TokenizedResponse] = Field(
- default=RunnerResponseType.TokenizedResponse, frozen=True
- )
+class TokenizedResponse(BaseRunnerResponse):
prompt_tokens: int
-class GenerationResponse(BaseRunnerResponse[RunnerResponseType.GenerationResponse]):
- type: Literal[RunnerResponseType.GenerationResponse] = Field(
- default=RunnerResponseType.GenerationResponse, frozen=True
- )
+class GenerationResponse(BaseRunnerResponse):
text: str
token: int
# logprobs: Optional[list[float]] = None # too big. we can change to be top-k
finish_reason: FinishReason | None = None
-class PrintResponse(BaseRunnerResponse[RunnerResponseType.PrintResponse]):
- type: Literal[RunnerResponseType.PrintResponse] = Field(
- default=RunnerResponseType.PrintResponse, frozen=True
- )
+class PrintResponse(BaseRunnerResponse):
text: str
-class FinishedResponse(BaseRunnerResponse[RunnerResponseType.FinishedResponse]):
- type: Literal[RunnerResponseType.FinishedResponse] = Field(
- default=RunnerResponseType.FinishedResponse, frozen=True
- )
+class FinishedResponse(BaseRunnerResponse):
+ pass
-class ErrorResponse(BaseRunnerResponse[RunnerResponseType.ErrorResponse]):
- type: Literal[RunnerResponseType.ErrorResponse] = Field(
- default=RunnerResponseType.ErrorResponse, frozen=True
- )
+class ErrorResponse(BaseRunnerResponse):
error_type: str
error_message: str
traceback: str
-RunnerResponse = Annotated[
+RunnerResponse = (
InitializedResponse
| TokenizedResponse
| GenerationResponse
| PrintResponse
| FinishedResponse
- | ErrorResponse,
- Field(discriminator="type"),
-]
-RunnerResponseTypeAdapter: TypeAdapter[RunnerResponse] = TypeAdapter(RunnerResponse)
+ | ErrorResponse
+)
diff --git a/src/exo/shared/types/worker/common.py b/src/exo/shared/types/worker/common.py
index 55441dd9..6dd29380 100644
--- a/src/exo/shared/types/worker/common.py
+++ b/src/exo/shared/types/worker/common.py
@@ -1,13 +1,13 @@
from enum import Enum
-from exo.shared.types.common import ID
+from exo.shared.types.common import Id
-class InstanceId(ID):
+class InstanceId(Id):
pass
-class RunnerId(ID):
+class RunnerId(Id):
pass
diff --git a/src/exo/shared/types/worker/communication.py b/src/exo/shared/types/worker/communication.py
index 0171acd6..7643af88 100644
--- a/src/exo/shared/types/worker/communication.py
+++ b/src/exo/shared/types/worker/communication.py
@@ -9,7 +9,6 @@ from exo.shared.types.worker.commands_runner import (
PrintResponse,
RunnerMessage,
RunnerResponse,
- RunnerResponseType,
)
### Utils - Runner Prints
@@ -17,7 +16,6 @@ from exo.shared.types.worker.commands_runner import (
def runner_print(text: str) -> None:
obj = PrintResponse(
- type=RunnerResponseType.PrintResponse,
text=text,
)
@@ -27,7 +25,6 @@ def runner_print(text: str) -> None:
def runner_write_error(error: Exception) -> None:
error_response: ErrorResponse = ErrorResponse(
- type=RunnerResponseType.ErrorResponse,
error_type=type(error).__name__,
error_message=str(error),
traceback=traceback.format_exc(),
diff --git a/src/exo/shared/types/worker/downloads.py b/src/exo/shared/types/worker/downloads.py
index aa5ee576..5c58b3e4 100644
--- a/src/exo/shared/types/worker/downloads.py
+++ b/src/exo/shared/types/worker/downloads.py
@@ -1,15 +1,6 @@
-from enum import Enum
-from typing import (
- Annotated,
- Literal,
- Union,
-)
-
-from pydantic import Field
-
from exo.shared.types.common import NodeId
from exo.shared.types.memory import Memory
-from exo.utils.pydantic_ext import CamelCaseModel
+from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
class DownloadProgressData(CamelCaseModel):
@@ -17,50 +8,26 @@ class DownloadProgressData(CamelCaseModel):
downloaded_bytes: Memory
-class DownloadStatus(str, Enum):
- Pending = "Pending"
- Downloading = "Downloading"
- Completed = "Completed"
- Failed = "Failed"
-
-
-class BaseDownloadProgress[DownloadStatusT: DownloadStatus](CamelCaseModel):
+class BaseDownloadProgress(TaggedModel):
node_id: NodeId
- download_status: DownloadStatusT
-class DownloadPending(BaseDownloadProgress[DownloadStatus.Pending]):
- download_status: Literal[DownloadStatus.Pending] = Field(
- default=DownloadStatus.Pending
- )
+class DownloadPending(BaseDownloadProgress):
+ pass
-class DownloadCompleted(BaseDownloadProgress[DownloadStatus.Completed]):
- download_status: Literal[DownloadStatus.Completed] = Field(
- default=DownloadStatus.Completed
- )
+class DownloadCompleted(BaseDownloadProgress):
+ pass
-class DownloadFailed(BaseDownloadProgress[DownloadStatus.Failed]):
- download_status: Literal[DownloadStatus.Failed] = Field(
- default=DownloadStatus.Failed
- )
+class DownloadFailed(BaseDownloadProgress):
error_message: str
-class DownloadOngoing(BaseDownloadProgress[DownloadStatus.Downloading]):
- download_status: Literal[DownloadStatus.Downloading] = Field(
- default=DownloadStatus.Downloading
- )
+class DownloadOngoing(BaseDownloadProgress):
download_progress: DownloadProgressData
-DownloadProgress = Annotated[
- Union[
- DownloadPending,
- DownloadCompleted,
- DownloadFailed,
- DownloadOngoing,
- ],
- Field(discriminator="download_status"),
-]
+DownloadProgress = (
+ DownloadPending | DownloadCompleted | DownloadFailed | DownloadOngoing
+)
diff --git a/src/exo/shared/types/worker/instances.py b/src/exo/shared/types/worker/instances.py
index d44a0e54..bb275e42 100644
--- a/src/exo/shared/types/worker/instances.py
+++ b/src/exo/shared/types/worker/instances.py
@@ -1,20 +1,19 @@
from enum import Enum
-from pydantic import BaseModel
-
from exo.shared.types.common import Host
from exo.shared.types.worker.common import InstanceId
from exo.shared.types.worker.runners import (
ShardAssignments,
)
+from exo.utils.pydantic_ext import CamelCaseModel
class InstanceStatus(str, Enum):
- ACTIVE = "ACTIVE"
- INACTIVE = "INACTIVE"
+ Active = "Active"
+ Inactive = "Inactive"
-class Instance(BaseModel):
+class Instance(CamelCaseModel):
instance_id: InstanceId
instance_type: InstanceStatus
shard_assignments: ShardAssignments
diff --git a/src/exo/shared/types/worker/ops.py b/src/exo/shared/types/worker/ops.py
index 386e2f4b..a0ac696d 100644
--- a/src/exo/shared/types/worker/ops.py
+++ b/src/exo/shared/types/worker/ops.py
@@ -1,86 +1,49 @@
-from enum import Enum
-from typing import Annotated, Generic, Literal, TypeVar, Union
-
-from pydantic import BaseModel, Field
-
from exo.shared.types.common import Host
from exo.shared.types.events import InstanceId
from exo.shared.types.tasks import Task
from exo.shared.types.worker.common import RunnerId
from exo.shared.types.worker.shards import ShardMetadata
+from exo.utils.pydantic_ext import TaggedModel
-class RunnerOpType(str, Enum):
- ASSIGN_RUNNER = "assign_runner"
- UNASSIGN_RUNNER = "unassign_runner"
- RUNNER_UP = "runner_up"
- RUNNER_DOWN = "runner_down"
- RUNNER_FAILED = "runner_failed"
- CHAT_COMPLETION = "chat_completion"
-
-
-RunnerOpT = TypeVar("RunnerOpT", bound=RunnerOpType)
-
-
-class BaseRunnerOp(BaseModel, Generic[RunnerOpT]):
- op_type: RunnerOpT
+class BaseRunnerOp(TaggedModel):
+ pass
-class AssignRunnerOp(BaseRunnerOp[Literal[RunnerOpType.ASSIGN_RUNNER]]):
- op_type: Literal[RunnerOpType.ASSIGN_RUNNER] = Field(
- default=RunnerOpType.ASSIGN_RUNNER, frozen=True
- )
+class AssignRunnerOp(BaseRunnerOp):
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
- )
+class UnassignRunnerOp(BaseRunnerOp):
runner_id: RunnerId
-class RunnerUpOp(BaseRunnerOp[Literal[RunnerOpType.RUNNER_UP]]):
- op_type: Literal[RunnerOpType.RUNNER_UP] = Field(
- default=RunnerOpType.RUNNER_UP, frozen=True
- )
+class RunnerUpOp(BaseRunnerOp):
runner_id: RunnerId
-class RunnerDownOp(BaseRunnerOp[Literal[RunnerOpType.RUNNER_DOWN]]):
- op_type: Literal[RunnerOpType.RUNNER_DOWN] = Field(
- default=RunnerOpType.RUNNER_DOWN, frozen=True
- )
+class RunnerDownOp(BaseRunnerOp):
runner_id: RunnerId
-class RunnerFailedOp(BaseRunnerOp[Literal[RunnerOpType.RUNNER_FAILED]]):
- op_type: Literal[RunnerOpType.RUNNER_FAILED] = Field(
- default=RunnerOpType.RUNNER_FAILED, frozen=True
- )
+class RunnerFailedOp(BaseRunnerOp):
runner_id: RunnerId
-class ExecuteTaskOp(BaseRunnerOp[Literal[RunnerOpType.CHAT_COMPLETION]]):
- op_type: Literal[RunnerOpType.CHAT_COMPLETION] = Field(
- default=RunnerOpType.CHAT_COMPLETION, frozen=True
- )
+class ExecuteTaskOp(BaseRunnerOp):
runner_id: RunnerId
task: Task
# Aggregate all runner operations into a single, strictly-typed union for dispatching.
-RunnerOp = Annotated[
- Union[
- AssignRunnerOp,
- UnassignRunnerOp,
- RunnerUpOp,
- RunnerDownOp,
- RunnerFailedOp,
- ExecuteTaskOp,
- ],
- Field(discriminator="op_type"),
-]
+RunnerOp = (
+ AssignRunnerOp
+ | UnassignRunnerOp
+ | RunnerUpOp
+ | RunnerDownOp
+ | RunnerFailedOp
+ | ExecuteTaskOp
+)
diff --git a/src/exo/shared/types/worker/runners.py b/src/exo/shared/types/worker/runners.py
index 2a1e75da..1a36a268 100644
--- a/src/exo/shared/types/worker/runners.py
+++ b/src/exo/shared/types/worker/runners.py
@@ -1,80 +1,54 @@
from collections.abc import Mapping
-from enum import Enum
-from typing import Annotated, Literal
-from pydantic import BaseModel, Field, TypeAdapter, model_validator
+from pydantic import model_validator
from exo.shared.types.common import NodeId
from exo.shared.types.models import ModelId
from exo.shared.types.worker.common import RunnerId
from exo.shared.types.worker.downloads import DownloadProgress
from exo.shared.types.worker.shards import ShardMetadata
+from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
-class RunnerStatusType(str, Enum):
- Downloading = "Downloading"
- Inactive = "Inactive"
- Starting = "Starting"
- Loaded = "Loaded"
- Running = "Running"
- Failed = "Failed"
+class BaseRunnerStatus(TaggedModel):
+ pass
-class BaseRunnerStatus[T: RunnerStatusType](BaseModel):
- runner_status: T
-
-
-class DownloadingRunnerStatus(BaseRunnerStatus[RunnerStatusType.Downloading]):
- runner_status: Literal[RunnerStatusType.Downloading] = Field(
- default=RunnerStatusType.Downloading
- )
+class DownloadingRunnerStatus(BaseRunnerStatus):
download_progress: DownloadProgress
-class InactiveRunnerStatus(BaseRunnerStatus[RunnerStatusType.Inactive]):
- runner_status: Literal[RunnerStatusType.Inactive] = Field(
- default=RunnerStatusType.Inactive
- )
+class InactiveRunnerStatus(BaseRunnerStatus):
+ pass
-class StartingRunnerStatus(BaseRunnerStatus[RunnerStatusType.Starting]):
- runner_status: Literal[RunnerStatusType.Starting] = Field(
- default=RunnerStatusType.Starting
- )
+class StartingRunnerStatus(BaseRunnerStatus):
+ pass
-class LoadedRunnerStatus(BaseRunnerStatus[RunnerStatusType.Loaded]):
- runner_status: Literal[RunnerStatusType.Loaded] = Field(
- default=RunnerStatusType.Loaded
- )
+class LoadedRunnerStatus(BaseRunnerStatus):
+ pass
-class RunningRunnerStatus(BaseRunnerStatus[RunnerStatusType.Running]):
- runner_status: Literal[RunnerStatusType.Running] = Field(
- default=RunnerStatusType.Running
- )
+class RunningRunnerStatus(BaseRunnerStatus):
+ pass
-class FailedRunnerStatus(BaseRunnerStatus[RunnerStatusType.Failed]):
- runner_status: Literal[RunnerStatusType.Failed] = Field(
- default=RunnerStatusType.Failed
- )
+class FailedRunnerStatus(BaseRunnerStatus):
error_message: str | None = None
-RunnerStatus = Annotated[
+RunnerStatus = (
DownloadingRunnerStatus
| InactiveRunnerStatus
| StartingRunnerStatus
| LoadedRunnerStatus
| RunningRunnerStatus
- | FailedRunnerStatus,
- Field,
-]
-RunnerStatusParser: TypeAdapter[RunnerStatus] = TypeAdapter(RunnerStatus)
+ | FailedRunnerStatus
+)
-class ShardAssignments(BaseModel):
+class ShardAssignments(CamelCaseModel):
model_id: ModelId
runner_to_shard: Mapping[RunnerId, ShardMetadata]
node_to_runner: Mapping[NodeId, RunnerId]
diff --git a/src/exo/shared/types/worker/shards.py b/src/exo/shared/types/worker/shards.py
index d0602877..887530cd 100644
--- a/src/exo/shared/types/worker/shards.py
+++ b/src/exo/shared/types/worker/shards.py
@@ -1,39 +1,26 @@
-from enum import Enum
-from typing import Annotated, Generic, Literal, Optional, TypeVar
+from pydantic import Field
-from pydantic import BaseModel, Field, TypeAdapter
+from exo.shared.types.models import ModelMetadata
+from exo.utils.pydantic_ext import TaggedModel
-from exo.shared.types.common import NodeId
-from exo.shared.types.models import ModelId, ModelMetadata
-
-class PartitionStrategy(str, Enum):
- pipeline = "pipeline"
-
-
-PartitionStrategyT = TypeVar(
- "PartitionStrategyT", bound=PartitionStrategy, covariant=True
-)
-
-
-class BaseShardMetadata(BaseModel, Generic[PartitionStrategyT]):
+class BaseShardMetadata(TaggedModel):
"""
Defines a specific shard of the model that is ready to be run on a device.
Replaces previous `Shard` object.
"""
model_meta: ModelMetadata
- partition_strategy: PartitionStrategyT
device_rank: int
world_size: int
# Error handling; equivalent to monkey-patch, but we can't monkey-patch runner.py
# This is kinda annoying because it allocates memory in the ShardMetadata object. Can be rethought after Shanghai.
immediate_exception: bool = False
- should_timeout: Optional[float] = None
+ should_timeout: float | None = None
-class PipelineShardMetadata(BaseShardMetadata[Literal[PartitionStrategy.pipeline]]):
+class PipelineShardMetadata(BaseShardMetadata):
"""
Pipeline parallelism shard meta.
@@ -41,12 +28,9 @@ class PipelineShardMetadata(BaseShardMetadata[Literal[PartitionStrategy.pipeline
where start_layer is inclusive and end_layer is exclusive.
"""
- partition_strategy: Literal[PartitionStrategy.pipeline] = Field(
- default=PartitionStrategy.pipeline, frozen=True
- )
- start_layer: Annotated[int, Field(ge=0)]
- end_layer: Annotated[int, Field(ge=0)]
- n_layers: Annotated[int, Field(ge=0)]
+ start_layer: int = Field(ge=0)
+ end_layer: int = Field(ge=0)
+ n_layers: int = Field(ge=0)
@property
def is_first_layer(self) -> bool:
@@ -62,17 +46,4 @@ class PipelineShardMetadata(BaseShardMetadata[Literal[PartitionStrategy.pipeline
)
-ShardMetadata = Annotated[
- PipelineShardMetadata, Field(discriminator="partition_strategy")
-]
-ShardMetadataParser: TypeAdapter[ShardMetadata] = TypeAdapter(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, BaseShardMetadata[PartitionStrategyT]]
+ShardMetadata = PipelineShardMetadata
diff --git a/src/exo/utils/pydantic_ext.py b/src/exo/utils/pydantic_ext.py
index 1bbedea2..5600d386 100644
--- a/src/exo/utils/pydantic_ext.py
+++ b/src/exo/utils/pydantic_ext.py
@@ -1,5 +1,13 @@
-from pydantic import BaseModel, ConfigDict
+# pyright: reportAny=false, reportUnknownArgumentType=false, reportUnknownVariableType=false
+
+from typing import Any, Self
+
+from pydantic import BaseModel, ConfigDict, model_serializer, model_validator
from pydantic.alias_generators import to_camel
+from pydantic_core.core_schema import (
+ SerializerFunctionWrapHandler,
+ ValidatorFunctionWrapHandler,
+)
class CamelCaseModel(BaseModel):
@@ -12,5 +20,20 @@ class CamelCaseModel(BaseModel):
validate_by_name=True,
extra="forbid",
# I want to reenable this ASAP, but it's causing an issue with TaskStatus
- # strict=True,
+ strict=True,
)
+
+
+class TaggedModel(CamelCaseModel):
+ @model_serializer(mode="wrap")
+ def _serialize(self, handler: SerializerFunctionWrapHandler):
+ inner = handler(self)
+ return {self.__class__.__name__: inner}
+
+ @model_validator(mode="wrap")
+ @classmethod
+ def _validate(cls, v: Any, handler: ValidatorFunctionWrapHandler) -> Self:
+ if isinstance(v, dict) and len(v) == 1 and cls.__name__ in v:
+ return handler(v[cls.__name__])
+
+ return handler(v)
diff --git a/src/exo/utils/pydantic_tagged.py b/src/exo/utils/pydantic_tagged.py
deleted file mode 100644
index 3840e7dd..00000000
--- a/src/exo/utils/pydantic_tagged.py
+++ /dev/null
@@ -1,229 +0,0 @@
-# pyright: reportAny=false, reportPrivateUsage=false, reportUnusedParameter=false, reportUnknownMemberType=false
-
-from collections.abc import Callable
-from types import get_original_bases
-from typing import (
- Any,
- ClassVar,
- Self,
- Union,
- cast,
- get_args,
- get_origin,
-)
-
-import pydantic
-from bidict import bidict
-from pydantic import (
- BaseModel,
- Field,
- TypeAdapter,
- model_serializer,
- model_validator,
-)
-from pydantic_core import (
- PydanticCustomError,
-)
-
-
-def tagged_union[T: Tagged[Any]](
- type_map: dict[str, type],
-) -> Callable[[type[T]], type[T]]:
- def _decorator(cls: type[T]):
- # validate and process the types
- tagged_union_cls = _ensure_single_tagged_union_base(cls)
- adapter_dict = _ensure_tagged_union_generic_is_union(tagged_union_cls)
- type_bidict = _ensure_bijection_between_union_members_and_type_map(
- set(adapter_dict.keys()), type_map
- )
-
- # inject the adapter and type class variables
- cast(type[_TaggedImpl[Any]], cls)._type_bidict = type_bidict
- cast(type[_TaggedImpl[Any]], cls)._adapter_dict = adapter_dict
-
- return cls
-
- return _decorator
-
-
-class Tagged[C](BaseModel):
- """
- Utility for helping with serializing unions as adjacently tagged with Pydantic.
-
- By default, Pydantic uses internally tagged union ser/de BUT to play nicely with
- other cross-language ser/de tools, you need adjacently tagged unions, and Pydantic
- doesn't support those out of the box.
- SEE: https://serde.rs/enum-representations.html#adjacently-tagged
-
- This type is a Pydantic model in its own right and can be used on fields of other
- Pydantic models. It must be used in combination with `tagged_union` decorator to work.
-
- Example usage:
- ```python
- FoobarUnion = Union[Foo, Bar, Baz]
-
- @tagged_union({
- "Foo": Foo,
- "Bar": Bar,
- "Baz": Baz,
- })
- class TaggedFoobarUnion(Tagged[FoobarUnion]): ...
- ```
- """
-
- t: str = Field(frozen=True)
- """
- The tag corresponding to the type of the object in the union.
- """
-
- c: C = Field(frozen=True)
- """
- The actual content of the object of that type.
- """
-
- @classmethod
- def from_(cls, c: C) -> Self:
- t = cast(type[_TaggedImpl[C]], cls)._type_bidict.inv[type(c)]
- return cls(t=t, c=c)
-
- @model_serializer
- def _model_dump(self) -> dict[str, Any]:
- cls = type(cast(_TaggedImpl[C], self))
- adapter = cls._adapter_dict[cls._type_bidict[self.t]]
- return {
- "t": self.t,
- "c": adapter.dump_python(self.c),
- }
-
- @model_validator(mode="before")
- @classmethod
- def _model_validate_before(cls, data: Any) -> Any:
- cls = cast(type[_TaggedImpl[C]], cls)
-
- # check object shape & check "t" type is `str`
- if not isinstance(data, dict):
- raise PydanticCustomError(
- "dict_type", "Wrong object type: expected a dictionary type"
- )
- if "t" not in data or "c" not in data or len(data) != 2: # pyright: ignore[reportUnknownArgumentType]
- raise ValueError(
- "Wrong object shape: expected exactly {t: <tag>, c: <content>}"
- )
- if not isinstance(data["t"], str):
- raise PydanticCustomError(
- "string_type", 'Wrong field type: expected "t" to be `str`'
- )
-
- # grab tag & content keys + look up the type based on the tag
- t = data["t"]
- c = cast(Any, data["c"])
- ccls = cls._type_bidict.get(t)
- if ccls is None:
- raise PydanticCustomError(
- "union_tag_not_found",
- 'Wrong "t"-value: could not find tag within this discriminated union',
- )
- cadapter = cls._adapter_dict[ccls]
-
- return {
- "t": t,
- "c": cadapter.validate_python(c),
- }
-
- @model_validator(mode="after")
- def _model_validate_after(self) -> Self:
- cls = type(cast(_TaggedImpl[C], self))
- ccls = type(self.c)
-
- # sanity check for consistency
- t = cls._type_bidict.inv.get(ccls)
- if t is None:
- raise ValueError(
- 'Wrong "c"-value: could not find a tag corresponding to the type of this value'
- )
- if t != self.t:
- raise ValueError(
- 'Wrong "t"-value: the provided tag for this content\'s type mismatches the configured tag'
- )
-
- return self
-
-
-class _TaggedImpl[C](Tagged[C]):
- _type_bidict: ClassVar[bidict[str, type]]
- _adapter_dict: ClassVar[dict[type, TypeAdapter[Any]]]
-
-
-def _ensure_single_tagged_union_base(cls: type[Any]) -> type[Any]:
- bases = get_original_bases(cls)
-
- # count up all the bases (generic removed) and store last found one
- cnt = 0
- last = None
- for b in bases:
- if pydantic._internal._generics.get_origin(b) == Tagged: # pyright: ignore[reportAttributeAccessIssue]
- last = cast(type[Tagged[Any]], b)
- cnt += 1
-
- # sanity-check the bases
- if last is None:
- raise TypeError(f"Expected {Tagged!r} to be a base-class of {cls!r}")
- if cnt > 1:
- raise TypeError(
- f"Expected only one {Tagged!r} base-class of {cls!r}, but got {cnt}"
- )
-
- return last
-
-
-def _ensure_tagged_union_generic_is_union(
- cls: type[Any],
-) -> dict[type, TypeAdapter[Any]]:
- # extract type of the generic argument
- base_generics = cast(Any, pydantic._internal._generics.get_args(cls)) # pyright: ignore[reportAttributeAccessIssue]
- assert len(base_generics) == 1
- union_cls = base_generics[0]
-
- # ensure the generic is a union => extract the members
- union_origin = get_origin(union_cls)
- if union_origin != Union:
- raise TypeError(
- f"Expected {Tagged!r} base-class to have its generic be a {Union!r}, but got {union_cls!r}"
- )
- union_members = get_args(union_cls)
-
- # typecheck each of the members, creating a type<->adapter mapping
- adapter_dict: dict[type, TypeAdapter[Any]] = {}
- for m in union_members:
- if not isinstance(m, type):
- raise TypeError(f"Expected union member {m!r} to be a type")
- adapter_dict[m] = TypeAdapter(m)
-
- return adapter_dict
-
-
-def _ensure_bijection_between_union_members_and_type_map(
- members: set[type], type_map: dict[str, type]
-) -> bidict[str, type]:
- mapped_members = set(type_map.values())
-
- illegal_members = mapped_members - members
- for m in illegal_members:
- raise TypeError(
- f"Expected type-map member {m!r} to be member of the union, but is not"
- )
- missing_members = members - mapped_members
- for m in missing_members:
- raise TypeError(
- f"Expected type-map to include a tag for member {m!r}, but is missing"
- )
- assert mapped_members == members
-
- tag_sets = {m: {t for t in type_map if type_map[t] == m} for m in mapped_members}
- for m, ts in tag_sets.items():
- if len(ts) > 1:
- raise TypeError(
- f"Expected a single tag per member of the union, but found {ts} for member {m!r}"
- )
-
- return bidict(type_map)
diff --git a/src/exo/utils/tests/test_tagged.py b/src/exo/utils/tests/test_tagged.py
index b138dcac..6d417ed9 100644
--- a/src/exo/utils/tests/test_tagged.py
+++ b/src/exo/utils/tests/test_tagged.py
@@ -1,9 +1,8 @@
-from typing import Union
-
+import anyio
import pytest
from pydantic import BaseModel, TypeAdapter, ValidationError
-from exo.utils.pydantic_tagged import Tagged, tagged_union # ← CHANGE ME
+from exo.utils.pydantic_ext import TaggedModel
def test_plain_union_prefers_first_member_when_shapes_are_identical():
@@ -22,161 +21,230 @@ def test_plain_union_prefers_first_member_when_shapes_are_identical():
def test_tagged_union_serializes_and_deserializes_two_identical_shapes_correctly():
- class Foo1(BaseModel):
+ class Foo1(TaggedModel):
x: int
- class Foo2(BaseModel):
+ class Foo2(TaggedModel):
x: int
- foos = Union[Foo1, Foo2]
-
- @tagged_union({"Foo1": Foo1, "Foo2": Foo2})
- class TaggedFoos(Tagged[foos]):
- pass
+ t1 = Foo1(x=1)
+ assert t1.model_dump() == {"Foo1": {"x": 1}}
- # ---- serialize (via custom model_serializer) ----
- t1 = TaggedFoos.from_(Foo1(x=1))
- assert t1.model_dump() == {"t": "Foo1", "c": {"x": 1}}
-
- t2 = TaggedFoos.from_(Foo2(x=2))
- assert t2.model_dump() == {"t": "Foo2", "c": {"x": 2}}
+ t2 = Foo2(x=2)
+ assert t2.model_dump() == {"Foo2": {"x": 2}}
# ---- deserialize (TypeAdapter -> model_validator(before)) ----
- ta = TypeAdapter(TaggedFoos)
+ ta = TypeAdapter[Foo1 | Foo2](Foo1 | Foo2)
- out1 = ta.validate_python({"t": "Foo1", "c": {"x": 10}})
- assert isinstance(out1.c, Foo1) and out1.c.x == 10
+ out1 = ta.validate_python({"Foo1": {"x": 10}})
+ assert isinstance(out1, Foo1) and out1.x == 10
- out2 = ta.validate_python({"t": "Foo2", "c": {"x": 20}})
- assert isinstance(out2.c, Foo2) and out2.c.x == 20
+ out2 = ta.validate_python({"Foo2": {"x": 20}})
+ assert isinstance(out2, Foo2) and out2.x == 20
def test_tagged_union_rejects_unknown_tag():
- class Foo1(BaseModel):
+ class Foo1(TaggedModel):
x: int
- class Foo2(BaseModel):
+ class Foo2(TaggedModel):
x: int
- foos = Union[Foo1, Foo2]
+ ta = TypeAdapter[Foo1 | Foo2](Foo1 | Foo2)
+ with pytest.raises(ValidationError):
+ ta.validate_python({"NotARealTag": {"x": 0}})
+
+
+def test_two_tagged_classes_with_different_shapes_are_independent_and_not_cross_deserializable():
+ class A1(TaggedModel):
+ x: int
- @tagged_union({"Foo1": Foo1, "Foo2": Foo2})
- class TaggedFoos(Tagged[foos]):
- pass
+ class A2(TaggedModel):
+ name: str
+
+ class B1(TaggedModel):
+ name: str
+
+ class B2(TaggedModel):
+ active: bool
+
+ a_payload = A1(x=123).model_dump()
+ b_payload = B1(name="neo").model_dump()
+
+ assert a_payload == {"A1": {"x": 123}}
+ assert b_payload == {"B1": {"name": "neo"}}
+
+ ta_a = TypeAdapter[A1 | A2](A1 | A2)
+ ta_b = TypeAdapter[B1 | B2](B1 | B2)
- ta = TypeAdapter(TaggedFoos)
with pytest.raises(ValidationError):
- ta.validate_python({"t": "NotARealTag", "c": {"x": 0}})
+ ta_a.validate_python(b_payload)
+ with pytest.raises(ValidationError):
+ ta_b.validate_python(a_payload)
-def test_multiple_tagged_classes_do_not_override_each_others_mappings():
- """
- Creating a *new* Tagged[T] class must not mutate the previously defined one.
- This checks both the tag mapping and the per-class adapter dicts.
- """
- class Foo1(BaseModel):
- x: int
+class Inner(TaggedModel):
+ x: int
- class Foo2(BaseModel):
- x: int
- foos = Union[Foo1, Foo2]
+class Outer(TaggedModel):
+ inner: Inner
- @tagged_union({"One": Foo1, "Two": Foo2})
- class TaggedEN(Tagged[foos]):
- pass
- # Sanity: initial mapping/behavior
- obj_en_1 = TaggedEN.from_(Foo1(x=5))
- assert obj_en_1.t == "One"
- obj_en_2 = TaggedEN.from_(Foo2(x=6))
- assert obj_en_2.t == "Two"
+class Wrapper(TaggedModel):
+ outer: Outer
+ label: str
- # Define a second, different mapping
- @tagged_union({"Uno": Foo1, "Dos": Foo2})
- class TaggedES(Tagged[foos]):
- pass
- # The two classes should have *independent* mappings
- # (not the same object, and not equal content)
- assert TaggedEN._type_bidict is not TaggedES._type_bidict # pyright: ignore
- assert TaggedEN._type_bidict != TaggedES._type_bidict # pyright: ignore
+class Container(TaggedModel):
+ items: list[Inner]
+ nested: Wrapper
- # Their adapters dicts should also be distinct objects
- assert TaggedEN._adapter_dict is not TaggedES._adapter_dict # pyright: ignore
- # And both should cover the same set of member types
- assert set(TaggedEN._adapter_dict.keys()) == {Foo1, Foo2} # pyright: ignore
- assert set(TaggedES._adapter_dict.keys()) == {Foo1, Foo2} # pyright: ignore
- # Re-check that EN behavior has NOT changed after ES was created
- obj_en_1_again = TaggedEN.from_(Foo1(x=7))
- obj_en_2_again = TaggedEN.from_(Foo2(x=8))
- assert obj_en_1_again.t == "One"
- assert obj_en_2_again.t == "Two"
+def test_single_level_tagging():
+ inner = Inner(x=10)
+ dumped = inner.model_dump()
+ assert dumped == {"Inner": {"x": 10}}
- # ES behavior is per its *own* mapping
- obj_es_1 = TaggedES.from_(Foo1(x=9))
- obj_es_2 = TaggedES.from_(Foo2(x=10))
- assert obj_es_1.t == "Uno"
- assert obj_es_2.t == "Dos"
+ restored = Inner.model_validate(dumped)
+ assert isinstance(restored, Inner)
+ assert restored.x == 10
- # And deserialization respects each class's mapping independently
- ta_en = TypeAdapter(TaggedEN)
- ta_es = TypeAdapter(TaggedES)
- out_en = ta_en.validate_python({"t": "Two", "c": {"x": 123}})
- assert isinstance(out_en.c, Foo2) and out_en.c.x == 123
+def test_nested_externally_tagged_union_serializes_recursively():
+ outer = Outer(inner=Inner(x=42))
+ dumped = outer.model_dump()
- out_es = ta_es.validate_python({"t": "Dos", "c": {"x": 456}})
- assert isinstance(out_es.c, Foo2) and out_es.c.x == 456
+ assert dumped == {"Outer": {"inner": {"Inner": {"x": 42}}}}
+ restored = Outer.model_validate(dumped)
+ assert isinstance(restored.inner, Inner)
+ assert restored.inner.x == 42
-def test_two_tagged_classes_with_different_shapes_are_independent_and_not_cross_deserializable():
- class A1(BaseModel):
- x: int
- class A2(BaseModel):
- name: str
+def test_two_level_nested_tagging():
+ outer = Outer(inner=Inner(x=123))
+ dumped = outer.model_dump()
+ assert dumped == {"Outer": {"inner": {"Inner": {"x": 123}}}}
- union_a = Union[A1, A2]
+ restored = Outer.model_validate(dumped)
+ assert isinstance(restored.inner, Inner)
+ assert restored.inner.x == 123
- @tagged_union({"One": A1, "Two": A2})
- class TaggedA(Tagged[union_a]):
- pass
- class B1(BaseModel):
- name: str
+def test_three_level_nested_tagging():
+ wrapper = Wrapper(label="deep", outer=Outer(inner=Inner(x=7)))
+ dumped = wrapper.model_dump()
+ # 3-level structure, each with exactly one tag
+ assert dumped == {
+ "Wrapper": {
+ "label": "deep",
+ "outer": {"Outer": {"inner": {"Inner": {"x": 7}}}},
+ }
+ }
- class B2(BaseModel):
- active: bool
+ restored = Wrapper.model_validate(dumped)
+ assert isinstance(restored.outer.inner, Inner)
+ assert restored.outer.inner.x == 7
+ assert restored.label == "deep"
- union_b = Union[B1, B2]
- # Note: using the SAME tag strings intentionally to ensure mappings are per-class
- @tagged_union({"One": B1, "Two": B2})
- class TaggedB(Tagged[union_b]):
- pass
+def test_lists_and_mixed_nested_structures():
+ container = Container(
+ items=[Inner(x=1), Inner(x=2)],
+ nested=Wrapper(label="mix", outer=Outer(inner=Inner(x=9))),
+ )
+ dumped = container.model_dump()
- # --- Per-class state must be independent ---
- assert TaggedA._type_bidict is not TaggedB._type_bidict # pyright: ignore
- assert TaggedA._adapter_dict is not TaggedB._adapter_dict # pyright: ignore
- assert set(TaggedA._adapter_dict.keys()) == {A1, A2} # pyright: ignore
- assert set(TaggedB._adapter_dict.keys()) == {B1, B2} # pyright: ignore
+ assert dumped == {
+ "Container": {
+ "items": [
+ {"Inner": {"x": 1}},
+ {"Inner": {"x": 2}},
+ ],
+ "nested": {
+ "Wrapper": {
+ "label": "mix",
+ "outer": {"Outer": {"inner": {"Inner": {"x": 9}}}},
+ }
+ },
+ }
+ }
- # --- Round-trip for each class with overlapping tag strings ---
- a_payload = TaggedA.from_(A1(x=123)).model_dump()
- b_payload = TaggedB.from_(B1(name="neo")).model_dump()
+ restored = Container.model_validate(dumped)
+ assert isinstance(restored.nested.outer.inner, Inner)
+ assert [i.x for i in restored.items] == [1, 2]
- assert a_payload == {"t": "One", "c": {"x": 123}}
- assert b_payload == {"t": "One", "c": {"name": "neo"}}
- # --- Cross-deserialization must fail despite overlapping "t" values ---
- ta_a = TypeAdapter(TaggedA)
- ta_b = TypeAdapter(TaggedB)
+def test_no_double_tagging_on_repeated_calls():
+ """Ensure multiple model_dump calls don't stack tags."""
+ inner = Inner(x=11)
+ dumped1 = inner.model_dump()
+ dumped2 = inner.model_dump()
+ assert dumped1 == dumped2 == {"Inner": {"x": 11}}
- with pytest.raises(ValidationError):
- ta_a.validate_python(b_payload) # TaggedA expects {"x": ...} for tag "One"
+ outer = Outer(inner=inner)
+ d1 = outer.model_dump()
+ d2 = outer.model_dump()
+ assert d1 == d2 == {"Outer": {"inner": {"Inner": {"x": 11}}}}
- with pytest.raises(ValidationError):
- ta_b.validate_python(a_payload) # TaggedB expects {"name": ...} for tag "One"
+
+class L3A(TaggedModel):
+ x: int
+
+
+class L3B(TaggedModel):
+ x: int
+
+
+class L3C(TaggedModel):
+ x: int
+
+
+L3 = L3A | L3B | L3C
+
+
+class L2A(TaggedModel):
+ child: L3
+
+
+class L2B(TaggedModel):
+ child: L3
+
+
+class L2C(TaggedModel):
+ child: L3
+
+
+L2 = L2A | L2B | L2C
+
+
+class L1A(TaggedModel):
+ child: L2
+
+
+class L1B(TaggedModel):
+ child: L2
+
+
+class L1C(TaggedModel):
+ child: L2
+
+
+L1 = L1A | L1B | L1C
+
+
+@pytest.mark.anyio
+async def test_tagged_union_is_fast():
+ # payload along the "C" path (worst case for DFS if branches are tried A->B->C)
+ payload = {"L1C": {"child": {"L2C": {"child": {"L3C": {"x": 123}}}}}}
+
+ with anyio.fail_after(0.1):
+ out = TypeAdapter(L1).validate_python(payload) # type: ignore
+
+ # Sanity check the result
+ assert out.__class__.__name__ == "L1C" # type: ignore
+ assert out.child.__class__.__name__ == "L2C" # type: ignore
+ assert out.child.child.__class__.__name__ == "L3C" # type: ignore
+ assert out.child.child.x == 123 # type: ignore
diff --git a/src/exo/worker/download/download_utils.py b/src/exo/worker/download/download_utils.py
index b03e59eb..03551db9 100644
--- a/src/exo/worker/download/download_utils.py
+++ b/src/exo/worker/download/download_utils.py
@@ -12,7 +12,7 @@ from urllib.parse import urljoin
import aiofiles
import aiofiles.os as aios
import aiohttp
-from pydantic import BaseModel, DirectoryPath, Field, PositiveInt, TypeAdapter
+from pydantic import BaseModel, DirectoryPath, Field, PositiveInt, TypeAdapter, ConfigDict
from exo.shared.constants import EXO_HOME
from exo.shared.types.worker.shards import ShardMetadata
@@ -53,8 +53,7 @@ class RepoFileDownloadProgress(BaseModel):
status: Literal["not_started", "in_progress", "complete"]
start_time: float
- class Config:
- frozen = True
+ model_config = ConfigDict(frozen = True)
class RepoDownloadProgress(BaseModel):
@@ -88,8 +87,9 @@ class RepoDownloadProgress(BaseModel):
# fine-grained file progress keyed by file_path
file_progress: Dict[str, RepoFileDownloadProgress] = Field(default_factory=dict)
- class Config:
+ model_config = ConfigDict(
frozen = True # allow use as dict keys if desired
+ )
def build_model_path(model_id: str) -> DirectoryPath:
diff --git a/src/exo/worker/download/impl_shard_downloader.py b/src/exo/worker/download/impl_shard_downloader.py
index 6f49c3fb..67f3236c 100644
--- a/src/exo/worker/download/impl_shard_downloader.py
+++ b/src/exo/worker/download/impl_shard_downloader.py
@@ -5,7 +5,6 @@ from typing import AsyncIterator, Callable, Dict, List, Optional
from exo.shared.models.model_cards import MODEL_CARDS
from exo.shared.models.model_meta import get_model_meta
from exo.shared.types.worker.shards import (
- PartitionStrategy,
PipelineShardMetadata,
ShardMetadata,
)
@@ -24,7 +23,6 @@ async def build_base_shard(model_id: str) -> Optional[ShardMetadata]:
# print(f"build_base_shard {model_id=} {model_meta=}")
return PipelineShardMetadata(
model_meta=model_meta,
- partition_strategy=PartitionStrategy.pipeline,
device_rank=0,
world_size=1,
start_layer=0,
@@ -39,7 +37,6 @@ async def build_full_shard(model_id: str) -> Optional[PipelineShardMetadata]:
return None
return PipelineShardMetadata(
model_meta=base_shard.model_meta,
- partition_strategy=base_shard.partition_strategy,
device_rank=base_shard.device_rank,
world_size=base_shard.world_size,
start_layer=base_shard.start_layer,
diff --git a/src/exo/worker/download/shard_downloader.py b/src/exo/worker/download/shard_downloader.py
index 30615222..bd8ab417 100644
--- a/src/exo/worker/download/shard_downloader.py
+++ b/src/exo/worker/download/shard_downloader.py
@@ -6,7 +6,6 @@ from typing import AsyncIterator, Callable
from exo.shared.types.memory import Memory
from exo.shared.types.models import ModelId, ModelMetadata
from exo.shared.types.worker.shards import (
- PartitionStrategy,
PipelineShardMetadata,
ShardMetadata,
)
@@ -57,7 +56,6 @@ class ShardDownloader(ABC):
storage_size=Memory.from_bytes(0),
n_layers=1,
),
- partition_strategy=PartitionStrategy.pipeline,
device_rank=0,
world_size=1,
start_layer=0,
@@ -107,7 +105,6 @@ class NoopShardDownloader(ShardDownloader):
storage_size=Memory.from_bytes(0),
n_layers=1,
),
- partition_strategy=PartitionStrategy.pipeline,
device_rank=0,
world_size=1,
start_layer=0,
diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py
index 59cf2ca6..2dc09559 100644
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -12,7 +12,7 @@ from loguru import logger
from exo.routing.connection_message import ConnectionMessage, ConnectionMessageType
from exo.shared.apply import apply
-from exo.shared.types.commands import ForwarderCommand, RequestEventLog, TaggedCommand
+from exo.shared.types.commands import ForwarderCommand, RequestEventLog
from exo.shared.types.common import NodeId
from exo.shared.types.events import (
ChunkGenerated,
@@ -25,7 +25,6 @@ from exo.shared.types.events import (
NodePerformanceMeasured,
RunnerDeleted,
RunnerStatusUpdated,
- TaggedEvent,
TaskFailed,
TaskStateUpdated,
TopologyEdgeCreated,
@@ -50,7 +49,6 @@ from exo.shared.types.worker.ops import (
RunnerDownOp,
RunnerFailedOp,
RunnerOp,
- RunnerOpType,
RunnerUpOp,
UnassignRunnerOp,
)
@@ -120,19 +118,19 @@ class Worker:
),
)
+ async def memory_monitor_callback(
+ memory_profile: MemoryPerformanceProfile,
+ ) -> None:
+ await self.event_publisher(
+ NodeMemoryMeasured(node_id=self.node_id, memory=memory_profile)
+ )
+
# END CLEANUP
async with create_task_group() as tg:
self._tg = tg
tg.start_soon(start_polling_node_metrics, resource_monitor_callback)
- async def memory_monitor_callback(
- memory_profile: MemoryPerformanceProfile,
- ) -> None:
- await self.event_publisher(
- NodeMemoryMeasured(node_id=self.node_id, memory=memory_profile)
- )
-
tg.start_soon(start_polling_memory_metrics, memory_monitor_callback)
tg.start_soon(self._connection_message_event_writer)
tg.start_soon(self._resend_out_for_delivery)
@@ -154,8 +152,8 @@ class Worker:
async def _event_applier(self):
with self.global_event_receiver as events:
async for event in events:
- self.event_buffer.ingest(event.origin_idx, event.tagged_event.c)
- event_id = event.tagged_event.c.event_id
+ self.event_buffer.ingest(event.origin_idx, event.event)
+ event_id = event.event.event_id
if event_id in self.out_for_delivery:
del self.out_for_delivery[event_id]
@@ -256,9 +254,7 @@ class Worker:
await self.command_sender.send(
ForwarderCommand(
origin=self.node_id,
- tagged_command=TaggedCommand.from_(
- RequestEventLog(since_idx=0)
- ),
+ command=RequestEventLog(since_idx=0),
)
)
finally:
@@ -507,7 +503,7 @@ class Worker:
await queue.put(
TaskStateUpdated(
task_id=op.task.task_id,
- task_status=TaskStatus.RUNNING,
+ task_status=TaskStatus.Running,
)
)
@@ -528,14 +524,14 @@ class Worker:
)
if op.task.task_id in self.state.tasks:
- self.state.tasks[op.task.task_id].task_status = TaskStatus.COMPLETE
+ self.state.tasks[op.task.task_id].task_status = TaskStatus.Complete
if assigned_runner.shard_metadata.device_rank == 0:
# kind of hack - we don't want to wait for the round trip for this to complete
await queue.put(
TaskStateUpdated(
task_id=op.task.task_id,
- task_status=TaskStatus.COMPLETE,
+ task_status=TaskStatus.Complete,
)
)
@@ -582,18 +578,18 @@ class Worker:
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:
+ match op:
+ case AssignRunnerOp():
event_generator = self._execute_assign_op(op)
- case RunnerOpType.UNASSIGN_RUNNER:
+ case UnassignRunnerOp():
event_generator = self._execute_unassign_op(op)
- case RunnerOpType.RUNNER_UP:
+ case RunnerUpOp():
event_generator = self._execute_runner_up_op(op)
- case RunnerOpType.RUNNER_DOWN:
+ case RunnerDownOp():
event_generator = self._execute_runner_down_op(op)
- case RunnerOpType.RUNNER_FAILED:
+ case RunnerFailedOp():
event_generator = self._execute_runner_failed_op(op)
- case RunnerOpType.CHAT_COMPLETION:
+ case ExecuteTaskOp():
event_generator = self._execute_task_op(op)
async for event in event_generator:
@@ -624,7 +620,7 @@ class Worker:
if runner_id in self.assigned_runners:
yield TaskStateUpdated(
task_id=task_id,
- task_status=TaskStatus.FAILED,
+ task_status=TaskStatus.Failed,
)
yield TaskFailed(
@@ -638,7 +634,7 @@ class Worker:
fe = ForwarderEvent(
origin_idx=self.local_event_index,
origin=self.node_id,
- tagged_event=TaggedEvent.from_(event),
+ event=event,
)
await self.local_event_sender.send(fe)
self.out_for_delivery[event.event_id] = fe
diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py
index bf32f960..27dd5e75 100644
--- a/src/exo/worker/plan.py
+++ b/src/exo/worker/plan.py
@@ -6,7 +6,7 @@ from exo.shared.types.events import (
)
from exo.shared.types.tasks import Task, TaskId, TaskStatus
from exo.shared.types.worker.common import RunnerId
-from exo.shared.types.worker.downloads import DownloadStatus
+from exo.shared.types.worker.downloads import DownloadCompleted
from exo.shared.types.worker.instances import Instance, InstanceStatus
from exo.shared.types.worker.ops import (
AssignRunnerOp,
@@ -23,8 +23,8 @@ from exo.shared.types.worker.runners import (
InactiveRunnerStatus,
LoadedRunnerStatus,
RunnerStatus,
- RunnerStatusType,
RunningRunnerStatus,
+ StartingRunnerStatus,
)
from exo.worker.common import AssignedRunner
@@ -45,14 +45,12 @@ def unassign_runners(
# If our instance is in 'downloading' or 'assigned' state, then we know the runner is stale. These are part of AssignRunnerOp and should be blocking.
for assigned_runner_id in assigned_runners:
- if (
- assigned_runner_id in state_runners
- and isinstance(state_runners[assigned_runner_id], DownloadingRunnerStatus)
- # Not sure about this type ignore, i don't think it should be necessary
- and state_runners[assigned_runner_id].download_progress.download_status # type: ignore
- != DownloadStatus.Completed
- ):
- return UnassignRunnerOp(runner_id=assigned_runner_id)
+ if assigned_runner_id in state_runners:
+ status = state_runners[assigned_runner_id]
+ if isinstance(status, DownloadingRunnerStatus) and not isinstance(
+ status.download_progress, DownloadCompleted
+ ):
+ return UnassignRunnerOp(runner_id=assigned_runner_id)
return None
@@ -85,7 +83,7 @@ def spin_down_runners(
if (
runner_id in assigned_runners
and isinstance(assigned_runners[runner_id].status, LoadedRunnerStatus)
- and instance.instance_type == InstanceStatus.INACTIVE
+ and instance.instance_type == InstanceStatus.Inactive
):
return RunnerDownOp(runner_id=runner_id)
@@ -195,18 +193,19 @@ def spin_up_runners(
instance.shard_assignments.node_to_runner[worker_node_id]
].runner
is None
- and instance.instance_type == InstanceStatus.ACTIVE
+ and instance.instance_type == InstanceStatus.Active
):
# We are part of this instance, we want it up but it hasn't been spun up yet.
# Need to assert all other runners are ready before we can spin up.
ready_to_spin = True
for runner_id in instance.shard_assignments.node_to_runner.values():
- if runner_id in state_runners and state_runners[
- runner_id
- ].runner_status not in [
- RunnerStatusType.Inactive,
- RunnerStatusType.Starting,
- ]:
+ if runner_id in state_runners and isinstance(
+ state_runners[runner_id],
+ (
+ InactiveRunnerStatus,
+ StartingRunnerStatus,
+ ),
+ ):
ready_to_spin = False
if ready_to_spin:
@@ -229,13 +228,12 @@ def execute_task_op(
continue
assert runner_id in assigned_runners
runner = assigned_runners[runner_id]
- if runner.status.runner_status != RunnerStatusType.Loaded:
+ if not isinstance(runner.status, LoadedRunnerStatus):
continue # The only previous state to get to Running is from Loaded
for _, task in tasks.items():
if task.instance_id == instance_id and (
- task.task_status == TaskStatus.PENDING
- or task.task_status == TaskStatus.FAILED
+ task.task_status in (TaskStatus.Pending, TaskStatus.Failed)
):
if (
runner.shard_metadata.device_rank >= 1
diff --git a/src/exo/worker/tests/conftest.py b/src/exo/worker/tests/conftest.py
index 3c876418..3385866f 100644
--- a/src/exo/worker/tests/conftest.py
+++ b/src/exo/worker/tests/conftest.py
@@ -10,7 +10,6 @@ from exo.shared.types.tasks import (
ChatCompletionTask,
TaskId,
TaskStatus,
- TaskType,
)
from exo.shared.types.worker.common import InstanceId
from exo.shared.types.worker.instances import Instance, InstanceStatus
@@ -131,7 +130,7 @@ def instance(
return Instance(
instance_id=resolved_instance_id,
- instance_type=InstanceStatus.ACTIVE,
+ instance_type=InstanceStatus.Active,
shard_assignments=shard_assignments,
hosts=hosts(1),
)
@@ -161,8 +160,7 @@ def chat_completion_task(completion_create_params: ChatCompletionTaskParams):
task_id=resolved_task_id,
command_id=COMMAND_1_ID,
instance_id=resolved_instance_id,
- task_type=TaskType.CHAT_COMPLETION,
- task_status=TaskStatus.PENDING,
+ task_status=TaskStatus.Pending,
task_params=completion_create_params,
)
diff --git a/src/exo/worker/tests/test_handlers/test_handlers_happy.py b/src/exo/worker/tests/test_handlers/test_handlers_happy.py
index 86eb6ebf..89e1bc10 100644
--- a/src/exo/worker/tests/test_handlers/test_handlers_happy.py
+++ b/src/exo/worker/tests/test_handlers/test_handlers_happy.py
@@ -145,10 +145,10 @@ async def test_execute_task_op(
assert isinstance(events[0].runner_status, RunningRunnerStatus)
assert isinstance(events[1], TaskStateUpdated)
- assert events[1].task_status == TaskStatus.RUNNING # It tried to start.
+ assert events[1].task_status == TaskStatus.Running # It tried to start.
assert isinstance(events[-2], TaskStateUpdated)
- assert events[-2].task_status == TaskStatus.COMPLETE # It tried to start.
+ assert events[-2].task_status == TaskStatus.Complete # It tried to start.
assert isinstance(events[-1], RunnerStatusUpdated)
assert isinstance(
diff --git a/src/exo/worker/tests/test_integration/test_inference.py b/src/exo/worker/tests/test_integration/test_inference.py
index 4118896f..7b9b07d0 100644
--- a/src/exo/worker/tests/test_integration/test_inference.py
+++ b/src/exo/worker/tests/test_integration/test_inference.py
@@ -17,7 +17,6 @@ from exo.shared.types.tasks import (
Task,
TaskId,
TaskStatus,
- TaskType,
)
from exo.shared.types.worker.common import InstanceId, RunnerId
from exo.shared.types.worker.instances import (
@@ -57,7 +56,7 @@ async def test_runner_inference(
async with create_task_group() as tg:
tg.start_soon(worker.run)
instance_value: Instance = instance(INSTANCE_1_ID, NODE_A, RUNNER_1_ID)
- instance_value.instance_type = InstanceStatus.ACTIVE
+ instance_value.instance_type = InstanceStatus.Active
task: Task = chat_completion_task(INSTANCE_1_ID, TASK_1_ID)
await global_events.append_events(
@@ -120,7 +119,7 @@ async def test_2_runner_inference(
instance = Instance(
instance_id=INSTANCE_1_ID,
- instance_type=InstanceStatus.ACTIVE,
+ instance_type=InstanceStatus.Active,
shard_assignments=shard_assignments,
hosts=hosts(2),
)
@@ -190,7 +189,7 @@ async def test_2_runner_multi_message(
instance = Instance(
instance_id=INSTANCE_1_ID,
- instance_type=InstanceStatus.ACTIVE,
+ instance_type=InstanceStatus.Active,
shard_assignments=shard_assignments,
hosts=hosts(2),
)
@@ -218,8 +217,7 @@ async def test_2_runner_multi_message(
task_id=TASK_1_ID,
command_id=CommandId(),
instance_id=INSTANCE_1_ID,
- task_type=TaskType.CHAT_COMPLETION,
- task_status=TaskStatus.PENDING,
+ task_status=TaskStatus.Pending,
task_params=completion_create_params,
)
diff --git a/src/exo/worker/tests/test_integration/test_inference_sad.py b/src/exo/worker/tests/test_integration/test_inference_sad.py
index 82916549..595adb22 100644
--- a/src/exo/worker/tests/test_integration/test_inference_sad.py
+++ b/src/exo/worker/tests/test_integration/test_inference_sad.py
@@ -58,7 +58,7 @@ async def test_stream_response_failed_always(
async with create_task_group() as tg:
tg.start_soon(worker.run)
instance_value: Instance = instance(INSTANCE_1_ID, NODE_A, RUNNER_1_ID)
- instance_value.instance_type = InstanceStatus.ACTIVE
+ instance_value.instance_type = InstanceStatus.Active
async def mock_stream_response(
self: RunnerSupervisor,
@@ -88,8 +88,8 @@ async def test_stream_response_failed_always(
[
x
for x in events
- if isinstance(x.tagged_event.c, RunnerStatusUpdated)
- and isinstance(x.tagged_event.c.runner_status, FailedRunnerStatus)
+ if isinstance(x.event, RunnerStatusUpdated)
+ and isinstance(x.event.runner_status, FailedRunnerStatus)
]
)
== 3
@@ -99,13 +99,13 @@ async def test_stream_response_failed_always(
[
x
for x in events
- if isinstance(x.tagged_event.c, TaskStateUpdated)
- and x.tagged_event.c.task_status == TaskStatus.FAILED
+ if isinstance(x.event, TaskStateUpdated)
+ and x.event.task_status == TaskStatus.Failed
]
)
== 3
)
- assert any([isinstance(x.tagged_event.c, InstanceDeleted) for x in events])
+ assert any([isinstance(x.event, InstanceDeleted) for x in events])
await global_events.append_events(
[
@@ -152,7 +152,7 @@ async def test_stream_response_failed_once(
async with create_task_group() as tg:
tg.start_soon(worker.run)
instance_value: Instance = instance(INSTANCE_1_ID, NODE_A, RUNNER_1_ID)
- instance_value.instance_type = InstanceStatus.ACTIVE
+ instance_value.instance_type = InstanceStatus.Active
task: Task = chat_completion_task(INSTANCE_1_ID, TASK_1_ID)
await global_events.append_events(
@@ -186,8 +186,8 @@ async def test_stream_response_failed_once(
[
x
for x in events
- if isinstance(x.tagged_event.c, RunnerStatusUpdated)
- and isinstance(x.tagged_event.c.runner_status, FailedRunnerStatus)
+ if isinstance(x.event, RunnerStatusUpdated)
+ and isinstance(x.event.runner_status, FailedRunnerStatus)
]
)
== 1
@@ -197,8 +197,8 @@ async def test_stream_response_failed_once(
[
x
for x in events
- if isinstance(x.tagged_event.c, TaskStateUpdated)
- and x.tagged_event.c.task_status == TaskStatus.FAILED
+ if isinstance(x.event, TaskStateUpdated)
+ and x.event.task_status == TaskStatus.Failed
]
)
== 1
@@ -209,11 +209,11 @@ async def test_stream_response_failed_once(
seen_task_started, seen_task_finished = False, False
for wrapped_event in events:
- event = wrapped_event.tagged_event.c
+ event = wrapped_event.event
if isinstance(event, TaskStateUpdated):
- if event.task_status == TaskStatus.RUNNING:
+ if event.task_status == TaskStatus.Running:
seen_task_started = True
- if event.task_status == TaskStatus.COMPLETE:
+ if event.task_status == TaskStatus.Complete:
seen_task_finished = True
if isinstance(event, ChunkGenerated):
@@ -246,7 +246,7 @@ async def test_stream_response_timeout(
async with create_task_group() as tg:
tg.start_soon(worker.run)
instance_value: Instance = instance(INSTANCE_1_ID, NODE_A, RUNNER_1_ID)
- instance_value.instance_type = InstanceStatus.ACTIVE
+ instance_value.instance_type = InstanceStatus.Active
task: Task = chat_completion_task(INSTANCE_1_ID, TASK_1_ID)
task.task_params.messages[0].content = "EXO RUNNER MUST TIMEOUT"
@@ -269,8 +269,8 @@ async def test_stream_response_timeout(
[
x
for x in events
- if isinstance(x.tagged_event.c, RunnerStatusUpdated)
- and isinstance(x.tagged_event.c.runner_status, FailedRunnerStatus)
+ if isinstance(x.event, RunnerStatusUpdated)
+ and isinstance(x.event.runner_status, FailedRunnerStatus)
]
)
== 3
@@ -280,8 +280,8 @@ async def test_stream_response_timeout(
[
x
for x in events
- if isinstance(x.tagged_event.c, TaskStateUpdated)
- and x.tagged_event.c.task_status == TaskStatus.FAILED
+ if isinstance(x.event, TaskStateUpdated)
+ and x.event.task_status == TaskStatus.Failed
]
)
== 3
@@ -291,8 +291,8 @@ async def test_stream_response_timeout(
[
x
for x in events
- if isinstance(x.tagged_event.c, TaskFailed)
- and "timeouterror" in x.tagged_event.c.error_type.lower()
+ if isinstance(x.event, TaskFailed)
+ and "timeouterror" in x.event.error_type.lower()
]
)
== 3
diff --git a/src/exo/worker/tests/test_integration/test_instantiation.py b/src/exo/worker/tests/test_integration/test_instantiation.py
index fdba8ba1..4d852123 100644
--- a/src/exo/worker/tests/test_integration/test_instantiation.py
+++ b/src/exo/worker/tests/test_integration/test_instantiation.py
@@ -37,7 +37,7 @@ async def test_runner_spinup_timeout(
async with create_task_group() as tg:
tg.start_soon(worker.run)
instance_value: Instance = instance(INSTANCE_1_ID, NODE_A, RUNNER_1_ID)
- instance_value.instance_type = InstanceStatus.ACTIVE
+ instance_value.instance_type = InstanceStatus.Active
instance_value.shard_assignments.runner_to_shard[
RUNNER_1_ID
].should_timeout = 10
@@ -61,11 +61,11 @@ async def test_runner_spinup_timeout(
[
x
for x in events
- if isinstance(x.tagged_event.c, RunnerStatusUpdated)
- and isinstance(x.tagged_event.c.runner_status, FailedRunnerStatus)
+ if isinstance(x.event, RunnerStatusUpdated)
+ and isinstance(x.event.runner_status, FailedRunnerStatus)
]
)
== 3
)
- assert any([isinstance(x.tagged_event.c, InstanceDeleted) for x in events])
+ assert any([isinstance(x.event, InstanceDeleted) for x in events])
worker.shutdown()
diff --git a/src/exo/worker/tests/test_integration/test_instantiation_sad.py b/src/exo/worker/tests/test_integration/test_instantiation_sad.py
index f96c227f..e734ed49 100644
--- a/src/exo/worker/tests/test_integration/test_instantiation_sad.py
+++ b/src/exo/worker/tests/test_integration/test_instantiation_sad.py
@@ -38,7 +38,7 @@ async def test_runner_spinup_exception(
async with create_task_group() as tg:
tg.start_soon(worker.run)
instance_value: Instance = instance(INSTANCE_1_ID, NODE_A, RUNNER_1_ID)
- instance_value.instance_type = InstanceStatus.ACTIVE
+ instance_value.instance_type = InstanceStatus.Active
instance_value.shard_assignments.runner_to_shard[
RUNNER_1_ID
].immediate_exception = True
@@ -57,13 +57,13 @@ async def test_runner_spinup_exception(
[
x
for x in events
- if isinstance(x.tagged_event.c, RunnerStatusUpdated)
- and isinstance(x.tagged_event.c.runner_status, FailedRunnerStatus)
+ if isinstance(x.event, RunnerStatusUpdated)
+ and isinstance(x.event.runner_status, FailedRunnerStatus)
]
)
== 3
)
- assert any([isinstance(x.tagged_event.c, InstanceDeleted) for x in events])
+ assert any([isinstance(x.event, InstanceDeleted) for x in events])
worker.shutdown()
@@ -75,7 +75,7 @@ async def test_runner_spinup_timeout(
async with create_task_group() as tg:
tg.start_soon(worker.run)
instance_value: Instance = instance(INSTANCE_1_ID, NODE_A, RUNNER_1_ID)
- instance_value.instance_type = InstanceStatus.ACTIVE
+ instance_value.instance_type = InstanceStatus.Active
instance_value.shard_assignments.runner_to_shard[
RUNNER_1_ID
].should_timeout = 10
@@ -99,11 +99,11 @@ async def test_runner_spinup_timeout(
[
x
for x in events
- if isinstance(x.tagged_event.c, RunnerStatusUpdated)
- and isinstance(x.tagged_event.c.runner_status, FailedRunnerStatus)
+ if isinstance(x.event, RunnerStatusUpdated)
+ and isinstance(x.event.runner_status, FailedRunnerStatus)
]
)
== 3
)
- assert any([isinstance(x.tagged_event.c, InstanceDeleted) for x in events])
+ assert any([isinstance(x.event, InstanceDeleted) for x in events])
worker.shutdown()
diff --git a/src/exo/worker/tests/test_multimodel/test_inference_llama70B.py b/src/exo/worker/tests/test_multimodel/test_inference_llama70B.py
index 9ce8746f..60501d9c 100644
--- a/src/exo/worker/tests/test_multimodel/test_inference_llama70B.py
+++ b/src/exo/worker/tests/test_multimodel/test_inference_llama70B.py
@@ -22,7 +22,6 @@ from exo.shared.types.tasks import (
Task,
TaskId,
TaskStatus,
- TaskType,
)
from exo.shared.types.worker.common import InstanceId
from exo.shared.types.worker.instances import (
@@ -107,7 +106,7 @@ async def test_ttft(
instance = Instance(
instance_id=INSTANCE_1_ID,
- instance_type=InstanceStatus.ACTIVE,
+ instance_type=InstanceStatus.Active,
shard_assignments=shard_assignments,
hosts=hosts(1),
)
@@ -139,8 +138,7 @@ async def test_ttft(
task_id=TASK_1_ID,
command_id=COMMAND_1_ID,
instance_id=INSTANCE_1_ID,
- task_type=TaskType.CHAT_COMPLETION,
- task_status=TaskStatus.PENDING,
+ task_status=TaskStatus.Pending,
task_params=task1_params,
)
@@ -157,7 +155,7 @@ async def test_ttft(
first_chunk_seen_1 = False
time_to_first_token_1: None | float = None
while not first_chunk_seen_1:
- event = (await global_events.receive()).tagged_event.c
+ event = (await global_events.receive()).event
if isinstance(event, ChunkGenerated) and hasattr(event, "chunk"):
first_chunk_time_1 = time.time()
time_to_first_token_1 = first_chunk_time_1 - task_created_time_1
@@ -192,8 +190,7 @@ async def test_ttft(
task_id=TASK_2_ID,
command_id=COMMAND_2_ID,
instance_id=INSTANCE_1_ID,
- task_type=TaskType.CHAT_COMPLETION,
- task_status=TaskStatus.PENDING,
+ task_status=TaskStatus.Pending,
task_params=task2_params,
)
@@ -211,7 +208,7 @@ async def test_ttft(
first_chunk_seen_2 = False
time_to_first_token_2: float | None = None
while not first_chunk_seen_2:
- event = (await global_events.receive()).tagged_event.c
+ event = (await global_events.receive()).event
if isinstance(event, ChunkGenerated) and hasattr(event, "chunk"):
first_chunk_time_2 = time.time()
time_to_first_token_2 = first_chunk_time_2 - task_created_time_2
@@ -344,7 +341,7 @@ async def test_2_runner_inference(
instance = Instance(
instance_id=INSTANCE_1_ID,
- instance_type=InstanceStatus.ACTIVE,
+ instance_type=InstanceStatus.Active,
shard_assignments=shard_assignments,
hosts=hosts(2),
)
@@ -424,7 +421,7 @@ async def test_parallel_inference(
instance = Instance(
instance_id=INSTANCE_1_ID,
- instance_type=InstanceStatus.ACTIVE,
+ instance_type=InstanceStatus.Active,
shard_assignments=shard_assignments,
hosts=hosts(2),
)
@@ -443,8 +440,7 @@ async def test_parallel_inference(
task_id=TASK_1_ID,
command_id=COMMAND_1_ID,
instance_id=INSTANCE_1_ID,
- task_type=TaskType.CHAT_COMPLETION,
- task_status=TaskStatus.PENDING,
+ task_status=TaskStatus.Pending,
task_params=completion_create_params_1,
)
@@ -462,8 +458,7 @@ async def test_parallel_inference(
task_id=TASK_2_ID,
command_id=COMMAND_2_ID,
instance_id=INSTANCE_1_ID,
- task_type=TaskType.CHAT_COMPLETION,
- task_status=TaskStatus.PENDING,
+ task_status=TaskStatus.Pending,
task_params=completion_create_params_2,
)
@@ -485,7 +480,7 @@ async def test_parallel_inference(
incomplete_task = (
TASK_2_ID
- if worker1.state.tasks[TASK_1_ID].task_status == TaskStatus.COMPLETE
+ if worker1.state.tasks[TASK_1_ID].task_status == TaskStatus.Complete
else TASK_2_ID
)
(
diff --git a/src/exo/worker/tests/test_plan/test_worker_plan.py b/src/exo/worker/tests/test_plan/test_worker_plan.py
index c04038a5..02f9612d 100644
--- a/src/exo/worker/tests/test_plan/test_worker_plan.py
+++ b/src/exo/worker/tests/test_plan/test_worker_plan.py
@@ -6,7 +6,6 @@ from exo.shared.types.tasks import (
ChatCompletionTask,
ChatCompletionTaskParams,
TaskStatus,
- TaskType,
)
from exo.shared.types.worker.common import WorkerStatus
from exo.shared.types.worker.downloads import (
@@ -85,7 +84,7 @@ def _get_test_cases() -> list[PlanTestCase]:
"downloaded": False,
}
],
- instance_status=InstanceStatus.INACTIVE,
+ instance_status=InstanceStatus.Inactive,
expected_op=UnassignRunnerOp(runner_id=RUNNER_1_ID),
),
make_test_case(
@@ -99,7 +98,7 @@ def _get_test_cases() -> list[PlanTestCase]:
"downloaded": True,
}
],
- instance_status=InstanceStatus.INACTIVE,
+ instance_status=InstanceStatus.Inactive,
expected_op=None,
),
PlanTestCase(
@@ -110,7 +109,7 @@ def _get_test_cases() -> list[PlanTestCase]:
INSTANCE_1_ID: [(RUNNER_1_ID, NODE_A, 0, InactiveRunnerStatus())]
},
model_id=MODEL_A_ID,
- instance_status=InstanceStatus.ACTIVE, # Either active or inactive should yield the same.
+ instance_status=InstanceStatus.Active, # Either active or inactive should yield the same.
),
expected_op=AssignRunnerOp(
instance_id=INSTANCE_1_ID,
@@ -153,7 +152,7 @@ def _get_test_cases() -> list[PlanTestCase]:
"downloaded": True,
}
],
- instance_status=InstanceStatus.ACTIVE,
+ instance_status=InstanceStatus.Active,
expected_op=RunnerUpOp(runner_id=RUNNER_1_ID),
),
make_test_case(
@@ -180,11 +179,11 @@ def _get_test_cases() -> list[PlanTestCase]:
{
"task_id": TASK_1_ID,
"instance_id": INSTANCE_1_ID,
- "status": TaskStatus.PENDING,
+ "status": TaskStatus.Pending,
"messages": [{"role": "user", "content": "Hello, world!"}],
}
],
- instance_status=InstanceStatus.ACTIVE,
+ instance_status=InstanceStatus.Active,
expected_op=None,
),
make_test_case(
@@ -209,11 +208,11 @@ def _get_test_cases() -> list[PlanTestCase]:
{
"task_id": TASK_1_ID,
"instance_id": INSTANCE_1_ID,
- "status": TaskStatus.PENDING,
+ "status": TaskStatus.Pending,
"messages": [{"role": "user", "content": "Hello, world!"}],
}
],
- instance_status=InstanceStatus.ACTIVE,
+ instance_status=InstanceStatus.Active,
expected_op=RunnerUpOp(runner_id=RUNNER_1_ID),
),
make_test_case(
@@ -227,7 +226,7 @@ def _get_test_cases() -> list[PlanTestCase]:
"downloaded": True,
}
],
- instance_status=InstanceStatus.INACTIVE,
+ instance_status=InstanceStatus.Inactive,
expected_op=RunnerDownOp(runner_id=RUNNER_1_ID),
),
make_test_case(
@@ -241,7 +240,7 @@ def _get_test_cases() -> list[PlanTestCase]:
"downloaded": True,
}
],
- instance_status=InstanceStatus.INACTIVE,
+ instance_status=InstanceStatus.Inactive,
expected_op=RunnerDownOp(runner_id=RUNNER_1_ID),
),
make_test_case(
@@ -259,19 +258,18 @@ def _get_test_cases() -> list[PlanTestCase]:
{
"task_id": TASK_1_ID,
"instance_id": INSTANCE_1_ID,
- "status": TaskStatus.PENDING,
+ "status": TaskStatus.Pending,
"messages": [{"role": "user", "content": "Hello, world!"}],
}
],
- instance_status=InstanceStatus.ACTIVE,
+ instance_status=InstanceStatus.Active,
expected_op=ExecuteTaskOp(
runner_id=RUNNER_1_ID,
task=ChatCompletionTask(
task_id=TASK_1_ID,
command_id=COMMAND_1_ID,
instance_id=INSTANCE_1_ID,
- task_type=TaskType.CHAT_COMPLETION,
- task_status=TaskStatus.PENDING,
+ task_status=TaskStatus.Pending,
task_params=ChatCompletionTaskParams(
model=str(MODEL_A_ID),
messages=[
@@ -304,11 +302,11 @@ def _get_test_cases() -> list[PlanTestCase]:
{
"task_id": TASK_1_ID,
"instance_id": INSTANCE_1_ID,
- "status": TaskStatus.PENDING,
+ "status": TaskStatus.Pending,
"messages": [{"role": "user", "content": "Hello, world!"}],
}
],
- instance_status=InstanceStatus.ACTIVE,
+ instance_status=InstanceStatus.Active,
expected_op=None,
),
make_test_case(
@@ -333,25 +331,24 @@ def _get_test_cases() -> list[PlanTestCase]:
{
"task_id": TASK_1_ID,
"instance_id": INSTANCE_1_ID,
- "status": TaskStatus.PENDING,
+ "status": TaskStatus.Pending,
"messages": [{"role": "user", "content": "Hello, world!"}],
}
],
- instance_status=InstanceStatus.ACTIVE,
+ instance_status=InstanceStatus.Active,
expected_op=ExecuteTaskOp(
runner_id=RUNNER_1_ID,
task=ChatCompletionTask(
task_id=TASK_1_ID,
command_id=COMMAND_1_ID,
instance_id=INSTANCE_1_ID,
- task_type=TaskType.CHAT_COMPLETION,
task_params=ChatCompletionTaskParams(
model=str(MODEL_A_ID),
messages=[
ChatCompletionMessage(role="user", content="Hello, world!")
],
),
- task_status=TaskStatus.PENDING,
+ task_status=TaskStatus.Pending,
),
),
),
@@ -377,25 +374,24 @@ def _get_test_cases() -> list[PlanTestCase]:
{
"task_id": TASK_1_ID,
"instance_id": INSTANCE_1_ID,
- "status": TaskStatus.PENDING,
+ "status": TaskStatus.Pending,
"messages": [{"role": "user", "content": "Hello, world!"}],
}
],
- instance_status=InstanceStatus.ACTIVE,
+ instance_status=InstanceStatus.Active,
expected_op=ExecuteTaskOp(
runner_id=RUNNER_1_ID,
task=ChatCompletionTask(
task_id=TASK_1_ID,
command_id=COMMAND_1_ID,
instance_id=INSTANCE_1_ID,
- task_type=TaskType.CHAT_COMPLETION,
task_params=ChatCompletionTaskParams(
model=str(MODEL_A_ID),
messages=[
ChatCompletionMessage(role="user", content="Hello, world!")
],
),
- task_status=TaskStatus.PENDING,
+ task_status=TaskStatus.Pending,
),
),
),
@@ -410,7 +406,7 @@ def _get_test_cases() -> list[PlanTestCase]:
"downloaded": True,
}
],
- instance_status=InstanceStatus.ACTIVE,
+ instance_status=InstanceStatus.Active,
expected_op=RunnerDownOp(runner_id=RUNNER_1_ID),
),
make_test_case(
@@ -431,7 +427,7 @@ def _get_test_cases() -> list[PlanTestCase]:
"downloaded": True,
},
],
- instance_status=InstanceStatus.ACTIVE,
+ instance_status=InstanceStatus.Active,
expected_op=RunnerDownOp(runner_id=RUNNER_1_ID),
),
make_test_case(
@@ -452,7 +448,7 @@ def _get_test_cases() -> list[PlanTestCase]:
"downloaded": True,
},
],
- instance_status=InstanceStatus.ACTIVE,
+ instance_status=InstanceStatus.Active,
expected_op=None,
),
make_test_case(
@@ -473,7 +469,7 @@ def _get_test_cases() -> list[PlanTestCase]:
"downloaded": True,
},
],
- instance_status=InstanceStatus.ACTIVE,
+ instance_status=InstanceStatus.Active,
expected_op=RunnerDownOp(runner_id=RUNNER_1_ID),
),
]
diff --git a/src/exo/worker/tests/test_plan/test_worker_plan_utils.py b/src/exo/worker/tests/test_plan/test_worker_plan_utils.py
index 4c7d12f9..f4681c11 100644
--- a/src/exo/worker/tests/test_plan/test_worker_plan_utils.py
+++ b/src/exo/worker/tests/test_plan/test_worker_plan_utils.py
@@ -9,7 +9,7 @@ from exo.shared.types.common import CommandId, NodeId
from exo.shared.types.memory import Memory
from exo.shared.types.models import ModelId, ModelMetadata
from exo.shared.types.state import State
-from exo.shared.types.tasks import ChatCompletionTask, TaskId, TaskStatus, TaskType
+from exo.shared.types.tasks import ChatCompletionTask, TaskId, TaskStatus
from exo.shared.types.worker.common import InstanceId, RunnerId, WorkerStatus
from exo.shared.types.worker.downloads import DownloadOngoing, DownloadProgressData
from exo.shared.types.worker.instances import Instance, InstanceStatus
@@ -146,7 +146,7 @@ def make_instance(
instance_id: InstanceId,
runner_specs: list[tuple[RunnerId, NodeId, int, RunnerStatus]],
model_id: ModelId = MODEL_A_ID,
- instance_status: InstanceStatus = InstanceStatus.ACTIVE,
+ instance_status: InstanceStatus = InstanceStatus.Active,
) -> tuple[Instance, dict[RunnerId, RunnerStatus], dict[NodeId, WorkerStatus]]:
"""Creates an instance with one or more runners."""
runner_to_shard: dict[RunnerId, PipelineShardMetadata] = {}
@@ -189,7 +189,7 @@ def make_state(
],
tasks: dict[TaskId, ChatCompletionTask] | None = None,
model_id: ModelId = MODEL_A_ID,
- instance_status: InstanceStatus = InstanceStatus.ACTIVE,
+ instance_status: InstanceStatus = InstanceStatus.Active,
) -> State:
"""Builds a full State from runner specs per instance, tasks, and defaults."""
if tasks is None:
@@ -224,7 +224,7 @@ def make_test_case(
tasks: list[TaskSpecDict] | None = None,
expected_op: Optional[RunnerOp] = None,
instance_id: InstanceId = INSTANCE_1_ID,
- instance_status: InstanceStatus = InstanceStatus.ACTIVE,
+ instance_status: InstanceStatus = InstanceStatus.Active,
model_id: ModelId = MODEL_A_ID,
command_id: CommandId = COMMAND_1_ID, # Default for tasks
) -> PlanTestCase:
@@ -244,8 +244,7 @@ def make_test_case(
instance_id=instance_id,
task_id=t["task_id"],
command_id=t.get("command_id", command_id),
- task_type=TaskType.CHAT_COMPLETION,
- task_status=t.get("status", TaskStatus.PENDING),
+ task_status=t.get("status", TaskStatus.Pending),
task_params=ChatCompletionTaskParams(
model=t.get("model", str(model_id)),
messages=[
diff --git a/src/exo/worker/tests/test_runner_connection.py b/src/exo/worker/tests/test_runner_connection.py
index 0eccf5d3..a887b866 100644
--- a/src/exo/worker/tests/test_runner_connection.py
+++ b/src/exo/worker/tests/test_runner_connection.py
@@ -72,7 +72,7 @@ async def check_runner_connection(
instance = Instance(
instance_id=INSTANCE_1_ID,
- instance_type=InstanceStatus.ACTIVE,
+ instance_type=InstanceStatus.Active,
shard_assignments=shard_assignments,
hosts=hosts(2),
)
diff --git a/src/exo/worker/tests/test_serdes.py b/src/exo/worker/tests/test_serdes.py
index bee86310..58c9c307 100644
--- a/src/exo/worker/tests/test_serdes.py
+++ b/src/exo/worker/tests/test_serdes.py
@@ -6,7 +6,7 @@ from exo.shared.types.common import Host
from exo.shared.types.tasks import Task, TaskId
from exo.shared.types.worker.commands_runner import (
ChatTaskMessage,
- RunnerMessageTypeAdapter,
+ RunnerMessage,
SetupMessage,
)
from exo.shared.types.worker.common import InstanceId
@@ -30,7 +30,7 @@ def test_supervisor_setup_message_serdes(
model_shard_meta=pipeline_shard_meta(1, 0),
hosts=hosts(1),
)
- assert_equal_serdes(setup_message, RunnerMessageTypeAdapter)
+ assert_equal_serdes(setup_message, TypeAdapter(RunnerMessage))
def test_supervisor_task_message_serdes(
@@ -40,4 +40,4 @@ def test_supervisor_task_message_serdes(
task_message = ChatTaskMessage(
task_data=task.task_params,
)
- assert_equal_serdes(task_message, RunnerMessageTypeAdapter)
+ assert_equal_serdes(task_message, TypeAdapter(RunnerMessage))
diff --git a/src/exo/worker/tests/test_supervisor/test_supervisor.py b/src/exo/worker/tests/test_supervisor/test_supervisor.py
index 6b44c9b9..9a03862c 100644
--- a/src/exo/worker/tests/test_supervisor/test_supervisor.py
+++ b/src/exo/worker/tests/test_supervisor/test_supervisor.py
@@ -7,10 +7,10 @@ from exo.shared.openai_compat import FinishReason
from exo.shared.types.chunks import TokenChunk
from exo.shared.types.common import Host
from exo.shared.types.tasks import (
+ ChatCompletionTask,
ChatCompletionTaskParams,
Task,
TaskId,
- TaskType,
)
from exo.shared.types.worker.common import InstanceId
from exo.shared.types.worker.shards import PipelineShardMetadata
@@ -143,7 +143,7 @@ async def test_supervisor_early_stopping(
task = chat_completion_task(instance_id, TaskId())
max_tokens = 50
- assert task.task_type == TaskType.CHAT_COMPLETION
+ assert isinstance(task, ChatCompletionTask)
print(f"chat_completion_task.task_params: {task.task_params}")
assert isinstance(task.task_params, ChatCompletionTaskParams)
task_params: ChatCompletionTaskParams = task.task_params
diff --git a/src/exo/worker/tests/worker_management.py b/src/exo/worker/tests/worker_management.py
index 34b6db13..ad7e346d 100644
--- a/src/exo/worker/tests/worker_management.py
+++ b/src/exo/worker/tests/worker_management.py
@@ -6,7 +6,7 @@ from anyio import fail_after
from exo.routing.topics import ConnectionMessage, ForwarderCommand, ForwarderEvent
from exo.shared.types.chunks import TokenChunk
from exo.shared.types.common import NodeId
-from exo.shared.types.events import ChunkGenerated, Event, TaggedEvent, TaskStateUpdated
+from exo.shared.types.events import ChunkGenerated, Event, TaskStateUpdated
from exo.shared.types.tasks import TaskId, TaskStatus
from exo.utils.channels import Receiver, Sender, channel
from exo.worker.download.shard_downloader import NoopShardDownloader, ShardDownloader
@@ -24,7 +24,7 @@ class WorkerMailbox:
await self.sender.send(
ForwarderEvent(
origin=origin,
- tagged_event=TaggedEvent.from_(event),
+ event=event,
origin_idx=self.counter,
)
)
@@ -105,7 +105,7 @@ async def read_streaming_response(
token_count = 0
extra_events: list[Event] = []
- event = (await global_event_receiver.receive()).tagged_event.c
+ event = (await global_event_receiver.receive()).event
extra_events.append(event)
from loguru import logger
@@ -116,17 +116,17 @@ async def read_streaming_response(
if filter_task:
while not (
isinstance(event, TaskStateUpdated)
- and event.task_status == TaskStatus.RUNNING
+ and event.task_status == TaskStatus.Running
and event.task_id == filter_task
):
- event = (await global_event_receiver.receive()).tagged_event.c
+ event = (await global_event_receiver.receive()).event
extra_events.append(event)
for event in extra_events:
if isinstance(event, TaskStateUpdated):
- if event.task_status == TaskStatus.RUNNING:
+ if event.task_status == TaskStatus.Running:
seen_task_started += 1
- if event.task_status == TaskStatus.COMPLETE:
+ if event.task_status == TaskStatus.Complete:
seen_task_finished += 1
if isinstance(event, ChunkGenerated) and isinstance(
event.chunk, TokenChunk
@@ -137,11 +137,11 @@ async def read_streaming_response(
finish_reason = event.chunk.finish_reason
while not seen_task_finished:
- event = (await global_event_receiver.receive()).tagged_event.c
+ event = (await global_event_receiver.receive()).event
if isinstance(event, TaskStateUpdated):
- if event.task_status == TaskStatus.RUNNING:
+ if event.task_status == TaskStatus.Running:
seen_task_started += 1
- if event.task_status == TaskStatus.COMPLETE:
+ if event.task_status == TaskStatus.Complete:
seen_task_finished += 1
if isinstance(event, ChunkGenerated) and isinstance(
event.chunk, TokenChunk
@@ -167,7 +167,7 @@ async def until_event_with_timeout[T](
with fail_after(timeout):
while times_seen < multiplicity:
- event = (await global_event_receiver.receive()).tagged_event.c
+ event = (await global_event_receiver.receive()).event
if isinstance(event, event_type):
print(f"Wow! We got a {event}")
print(
← 76ed8a51 typecheck on ubuntu with install-nix-action
·
back to Exo
·
fix a race condition f25689d9 →