← back to Exo
tidy: fix justfile, run.sh, run formatter
11f8b4ef33ce81a73d14887ed3168a380a661817 · 2025-08-21 18:44:53 +0100 · Evan Quiney
Files touched
M justfileM run.shM src/exo/engines/mlx/auto_parallel.pyM src/exo/engines/mlx/utils_mlx.pyM src/exo/main.pyM src/exo/master/api.pyM src/exo/master/election_callback.pyM src/exo/master/forwarder_supervisor.pyM src/exo/master/main.pyM src/exo/master/placement.pyM src/exo/master/tests/api_utils_test.pyM src/exo/master/tests/conftest.pyM src/exo/master/tests/test_api.pyM src/exo/master/tests/test_forwarder_supervisor.pyM src/exo/master/tests/test_master.pyM src/exo/master/tests/test_placement.pyM src/exo/master/tests/test_placement_utils.pyM src/exo/master/tests/test_topology.pyM src/exo/master/utils/placement_utils.pyM src/exo/shared/__init__.pyM src/exo/shared/apply/__init__.pyM src/exo/shared/apply/apply.pyM src/exo/shared/constants.pyM src/exo/shared/db/__init__.pyM src/exo/shared/db/sqlite/__init__.pyM src/exo/shared/db/sqlite/config.pyM src/exo/shared/db/sqlite/connector.pyM src/exo/shared/db/sqlite/event_log_manager.pyM src/exo/shared/db/sqlite/types.pyM src/exo/shared/models/model_cards.pyM src/exo/shared/models/model_meta.pyM src/exo/shared/tests/__init__.pyM src/exo/shared/tests/test_node_id_persistence.pyM src/exo/shared/tests/test_sqlite_connector.pyM src/exo/shared/topology.pyM src/exo/shared/types/api.pyM src/exo/shared/types/common.pyM src/exo/shared/types/events/_events.pyM src/exo/shared/types/events/chunks.pyM src/exo/shared/types/events/commands.pyM src/exo/shared/types/events/components.pyM src/exo/shared/types/graphs/pydantic.pyM src/exo/shared/types/multiaddr.pyM src/exo/shared/types/request.pyM src/exo/shared/types/state.pyM src/exo/shared/types/tasks.pyM src/exo/shared/types/topology.pyM src/exo/shared/types/worker/commands_runner.pyM src/exo/shared/types/worker/common.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/resource_monitor.pyM src/exo/shared/types/worker/runners.pyM src/exo/shared/types/worker/shards.pyM src/exo/shared/utils.pyM src/exo/worker/common.pyM src/exo/worker/download/conftest.pyM src/exo/worker/download/download_utils.pyM src/exo/worker/download/huggingface_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/runner/communication.pyM src/exo/worker/runner/runner.pyM src/exo/worker/runner/runner_supervisor.pyM src/exo/worker/runner/utils.pyM src/exo/worker/tests/__init__.pyM src/exo/worker/tests/conftest.pyM src/exo/worker/tests/constants.pyM src/exo/worker/tests/test_download.pyM src/exo/worker/tests/test_handlers/conftest.pyM src/exo/worker/tests/test_handlers/test_handlers_happy.pyM src/exo/worker/tests/test_handlers/test_handlers_sad.pyM src/exo/worker/tests/test_handlers/utils.pyM src/exo/worker/tests/test_integration/conftest.pyM src/exo/worker/tests/test_integration/integration_utils.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_spinup_timeout.pyM src/exo/worker/tests/test_supervisor/test_memory.pyM src/exo/worker/tests/test_supervisor/test_oom.pyM src/exo/worker/tests/test_supervisor/test_supervisor.pyM src/exo/worker/tests/test_supervisor/test_supervisor_sad.pyM src/exo/worker/utils/macmon/__init__.pyM src/exo/worker/utils/macmon/macmon.pyM src/exo/worker/utils/profile.pyM src/exo/worker/utils/system_info.pyM src/exo/worker/worker.py
Diff
commit 11f8b4ef33ce81a73d14887ed3168a380a661817
Author: Evan Quiney <evanev7@gmail.com>
Date: Thu Aug 21 18:44:53 2025 +0100
tidy: fix justfile, run.sh, run formatter
---
justfile | 11 +-
run.sh | 10 +-
src/exo/engines/mlx/auto_parallel.py | 1 -
src/exo/engines/mlx/utils_mlx.py | 63 +-
src/exo/main.py | 2 +-
src/exo/master/api.py | 108 ++-
src/exo/master/election_callback.py | 8 +-
src/exo/master/forwarder_supervisor.py | 72 +-
src/exo/master/main.py | 146 ++--
src/exo/master/placement.py | 68 +-
src/exo/master/tests/api_utils_test.py | 11 +-
src/exo/master/tests/conftest.py | 19 +-
src/exo/master/tests/test_api.py | 10 +-
src/exo/master/tests/test_forwarder_supervisor.py | 131 ++--
src/exo/master/tests/test_master.py | 102 ++-
src/exo/master/tests/test_placement.py | 76 +-
src/exo/master/tests/test_placement_utils.py | 150 ++--
src/exo/master/tests/test_topology.py | 125 +++-
src/exo/master/utils/placement_utils.py | 37 +-
src/exo/shared/__init__.py | 1 -
src/exo/shared/apply/__init__.py | 2 +-
src/exo/shared/apply/apply.py | 96 ++-
src/exo/shared/constants.py | 1 +
src/exo/shared/db/__init__.py | 2 +-
src/exo/shared/db/sqlite/__init__.py | 2 +-
src/exo/shared/db/sqlite/config.py | 7 +-
src/exo/shared/db/sqlite/connector.py | 260 ++++---
src/exo/shared/db/sqlite/event_log_manager.py | 68 +-
src/exo/shared/db/sqlite/types.py | 37 +-
src/exo/shared/models/model_cards.py | 443 ++++++-----
src/exo/shared/models/model_meta.py | 42 +-
src/exo/shared/tests/__init__.py | 2 +-
src/exo/shared/tests/test_node_id_persistence.py | 9 +-
src/exo/shared/tests/test_sqlite_connector.py | 360 ++++++---
src/exo/shared/topology.py | 49 +-
src/exo/shared/types/api.py | 7 +-
src/exo/shared/types/common.py | 5 +-
src/exo/shared/types/events/_events.py | 50 +-
src/exo/shared/types/events/chunks.py | 5 +-
src/exo/shared/types/events/commands.py | 7 +-
src/exo/shared/types/events/components.py | 3 +-
src/exo/shared/types/graphs/pydantic.py | 2 +-
src/exo/shared/types/multiaddr.py | 41 +-
src/exo/shared/types/request.py | 3 +
src/exo/shared/types/state.py | 1 +
src/exo/shared/types/tasks.py | 1 +
src/exo/shared/types/topology.py | 27 +-
src/exo/shared/types/worker/commands_runner.py | 7 +-
src/exo/shared/types/worker/common.py | 15 +-
src/exo/shared/types/worker/downloads.py | 12 +-
src/exo/shared/types/worker/instances.py | 1 +
src/exo/shared/types/worker/ops.py | 36 +-
src/exo/shared/types/worker/resource_monitor.py | 16 +-
src/exo/shared/types/worker/runners.py | 33 +-
src/exo/shared/types/worker/shards.py | 18 +-
src/exo/shared/utils.py | 107 +--
src/exo/worker/common.py | 2 +-
src/exo/worker/download/conftest.py | 4 +-
src/exo/worker/download/download_utils.py | 817 +++++++++++++--------
src/exo/worker/download/huggingface_utils.py | 157 ++--
src/exo/worker/download/impl_shard_downloader.py | 277 ++++---
src/exo/worker/download/shard_downloader.py | 179 ++---
src/exo/worker/main.py | 112 +--
src/exo/worker/plan.py | 205 ++++--
src/exo/worker/runner/communication.py | 6 +-
src/exo/worker/runner/runner.py | 32 +-
src/exo/worker/runner/runner_supervisor.py | 96 +--
src/exo/worker/runner/utils.py | 35 +-
src/exo/worker/tests/__init__.py | 1 -
src/exo/worker/tests/conftest.py | 39 +-
src/exo/worker/tests/constants.py | 6 +-
src/exo/worker/tests/test_download.py | 4 +-
src/exo/worker/tests/test_handlers/conftest.py | 27 +-
.../tests/test_handlers/test_handlers_happy.py | 61 +-
.../tests/test_handlers/test_handlers_sad.py | 39 +-
src/exo/worker/tests/test_handlers/utils.py | 3 +-
src/exo/worker/tests/test_integration/conftest.py | 18 +-
.../tests/test_integration/integration_utils.py | 28 +-
.../tests/test_integration/test_inference.py | 175 +++--
.../tests/test_integration/test_inference_sad.py | 178 +++--
.../tests/test_integration/test_instantiation.py | 68 +-
.../test_integration/test_instantiation_sad.py | 63 +-
.../test_multimodel/test_inference_llama70B.py | 213 ++++--
src/exo/worker/tests/test_plan/test_worker_plan.py | 477 ++++++------
.../tests/test_plan/test_worker_plan_utils.py | 77 +-
src/exo/worker/tests/test_runner_connection.py | 45 +-
src/exo/worker/tests/test_spinup_timeout.py | 18 +-
.../worker/tests/test_supervisor/test_memory.py | 12 +-
src/exo/worker/tests/test_supervisor/test_oom.py | 10 +-
.../tests/test_supervisor/test_supervisor.py | 26 +-
.../tests/test_supervisor/test_supervisor_sad.py | 11 +-
src/exo/worker/utils/macmon/__init__.py | 2 +-
src/exo/worker/utils/macmon/macmon.py | 5 +-
src/exo/worker/utils/profile.py | 31 +-
src/exo/worker/utils/system_info.py | 151 ++--
src/exo/worker/worker.py | 156 ++--
96 files changed, 4144 insertions(+), 2650 deletions(-)
diff --git a/justfile b/justfile
index 4265a568..53aaf70c 100644
--- a/justfile
+++ b/justfile
@@ -2,19 +2,16 @@ default:
@just --list
fmt:
- uv run ruff format master worker shared engines/*
+ uv run ruff format src
lint:
- uv run ruff check --fix master worker shared engines/*
+ uv run ruff check --fix src
lint-check:
- uv run ruff check master worker shared engines/*
+ uv run ruff check src
test:
- uv run pytest master worker shared engines/*
-
-test-fast:
- uv run pytest master shared engines/*
+ uv run pytest src
check:
uv run basedpyright --project pyproject.toml
diff --git a/run.sh b/run.sh
index 74e81181..4daf8186 100755
--- a/run.sh
+++ b/run.sh
@@ -36,14 +36,14 @@ fi
# First command (worker) - changes based on replica flag
if [ "$REPLICA" = true ]; then
- osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c bash -c 'export EXO_HOME=.exo_replica; uv run -m worker.main'\""
+ osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c bash -c 'export EXO_HOME=.exo_replica; uv run exo-worker'\""
else
- osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c uv run -m worker.main\""
+ osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c uv run exo-worker\""
fi
# Second command (master) - changes based on replica flag
if [ "$REPLICA" = true ]; then
- osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c bash -c 'export RUST_LOG=true EXO_RUN_AS_REPLICA=1 EXO_HOME=.exo_replica API_PORT=8001; uv run -m master.main'\""
+ osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c bash -c 'export RUST_LOG=true EXO_RUN_AS_REPLICA=1 EXO_HOME=.exo_replica API_PORT=8001; uv run exo-master'\""
else
- osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c bash -c 'export RUST_LOG=true; uv run -m master.main'\""
-fi
\ No newline at end of file
+ osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c bash -c 'export RUST_LOG=true; uv run exo-master'\""
+fi
diff --git a/src/exo/engines/mlx/auto_parallel.py b/src/exo/engines/mlx/auto_parallel.py
index 83598a7a..2e2589fa 100644
--- a/src/exo/engines/mlx/auto_parallel.py
+++ b/src/exo/engines/mlx/auto_parallel.py
@@ -2,7 +2,6 @@ from typing import Protocol, cast, override
import mlx.core as mx
import mlx.nn as nn
-
from exo.shared.types.worker.shards import PipelineShardMetadata
diff --git a/src/exo/engines/mlx/utils_mlx.py b/src/exo/engines/mlx/utils_mlx.py
index c21c8c92..60a21e30 100644
--- a/src/exo/engines/mlx/utils_mlx.py
+++ b/src/exo/engines/mlx/utils_mlx.py
@@ -24,20 +24,22 @@ from exo.worker.runner.communication import runner_print
# Needed for 8 bit model
resource.setrlimit(resource.RLIMIT_NOFILE, (2048, 4096))
+
def mx_barrier():
- mx.eval( # type: ignore
+ mx.eval( # type: ignore
mx.distributed.all_sum(
mx.array(1.0), stream=mx.default_stream(mx.Device(mx.cpu))
)
)
+
class HostList(RootModel[list[str]]):
@classmethod
def from_hosts(cls, hosts: list[Host]) -> "HostList":
return cls(root=[str(host) for host in hosts])
-def mlx_distributed_init(rank: int, hosts: list[Host]) -> mx.distributed.Group: # type: ignore
+def mlx_distributed_init(rank: int, hosts: list[Host]) -> mx.distributed.Group: # type: ignore
"""
Initialize the MLX distributed (runs in thread pool)
"""
@@ -79,18 +81,20 @@ def initialize_mlx(
return model, tokenizer, sampler
-def shard_and_load(model_shard_meta: ShardMetadata) -> tuple[nn.Module, TokenizerWrapper]:
- model_path = build_model_path(model_shard_meta.model_meta.model_id)
+def shard_and_load(
+ model_shard_meta: ShardMetadata,
+) -> tuple[nn.Module, TokenizerWrapper]:
+ model_path = build_model_path(model_shard_meta.model_meta.model_id)
runner_print(f"loading model from {model_path}")
- model, _ = load_model(model_path, lazy=True, strict=False) # type: ignore
+ model, _ = load_model(model_path, lazy=True, strict=False) # type: ignore
assert isinstance(model, nn.Module)
tokenizer = load_tokenizer(model_path)
assert isinstance(tokenizer, TokenizerWrapper)
model = auto_parallel(model, model_shard_meta)
- mx.eval(model.parameters()) # type: ignore
+ mx.eval(model.parameters()) # type: ignore
# Synchronize processes before generation to avoid timeout
mx_barrier()
@@ -112,22 +116,24 @@ async def apply_chat_template(
# Filter out None values, keeping relevant keys for the model
formatted_messages = []
for message in messages_dicts:
- filtered_message: dict[str, Any] = {k: v for k, v in message.items() if v is not None} # type: ignore
-
+ filtered_message: dict[str, Any] = {
+ k: v for k, v in message.items() if v is not None
+ } # type: ignore
+
# Verify we have required fields
if "role" not in filtered_message:
raise ValueError(f"Message missing 'role' field: {filtered_message}")
if "content" not in filtered_message and "thinking" not in filtered_message:
# If neither content nor thinking is present, skip this message
continue
-
- formatted_messages.append(filtered_message) # type: ignore
+
+ formatted_messages.append(filtered_message) # type: ignore
messages_dicts = formatted_messages
prompt: str = await loop.run_in_executor(
executor=mlx_executor,
- func=lambda: tokenizer.apply_chat_template( # type: ignore
+ func=lambda: tokenizer.apply_chat_template( # type: ignore
messages_dicts,
tokenize=False,
add_generation_prompt=True,
@@ -136,6 +142,7 @@ async def apply_chat_template(
return prompt
+
async def warmup_inference(
mlx_executor: concurrent.futures.ThreadPoolExecutor,
model: nn.Module,
@@ -143,7 +150,7 @@ async def warmup_inference(
sampler: Callable[[mx.array], mx.array],
) -> int:
loop = asyncio.get_running_loop()
-
+
warmup_prompt = await apply_chat_template(
mlx_executor=mlx_executor,
tokenizer=tokenizer,
@@ -151,15 +158,15 @@ async def warmup_inference(
model="warmup",
messages=[
ChatCompletionMessage(
- role='user',
- content='Prompt to warm up the inference engine. Repeat this.'
+ role="user",
+ content="Prompt to warm up the inference engine. Repeat this.",
)
- ]
+ ],
),
)
-
+
tokens_generated = 0
-
+
def _generate_warmup():
nonlocal tokens_generated
for _ in stream_generate(
@@ -170,10 +177,10 @@ async def warmup_inference(
sampler=sampler,
):
tokens_generated += 1
-
+
await loop.run_in_executor(mlx_executor, _generate_warmup)
mx_barrier()
-
+
return tokens_generated
@@ -181,12 +188,12 @@ def mlx_force_oom(size: int = 40000) -> None:
"""
Force an Out-Of-Memory (OOM) error in MLX by performing large tensor operations.
"""
- mx.set_default_device(mx.gpu) # type: ignore
- a = mx.random.uniform(shape=(size, size), dtype=mx.float32) # type: ignore
- b = mx.random.uniform(shape=(size, size), dtype=mx.float32) # type: ignore
- mx.eval(a, b) # type: ignore
- c = mx.matmul(a, b) # type: ignore
- d = mx.matmul(a, c) # type: ignore
- e = mx.matmul(b, c) # type: ignore
- f = mx.sigmoid(d + e) # type: ignore
- mx.eval(f) # type: ignore
+ mx.set_default_device(mx.gpu) # type: ignore
+ a = mx.random.uniform(shape=(size, size), dtype=mx.float32) # type: ignore
+ b = mx.random.uniform(shape=(size, size), dtype=mx.float32) # type: ignore
+ mx.eval(a, b) # type: ignore
+ c = mx.matmul(a, b) # type: ignore
+ d = mx.matmul(a, c) # type: ignore
+ e = mx.matmul(b, c) # type: ignore
+ f = mx.sigmoid(d + e) # type: ignore
+ mx.eval(f) # type: ignore
diff --git a/src/exo/main.py b/src/exo/main.py
index 708a6a64..46b4ca54 100644
--- a/src/exo/main.py
+++ b/src/exo/main.py
@@ -1,2 +1,2 @@
def main():
- print("Hello world!")
+ print("Hello world!")
diff --git a/src/exo/master/api.py b/src/exo/master/api.py
index 207983f3..a347f7d4 100644
--- a/src/exo/master/api.py
+++ b/src/exo/master/api.py
@@ -1,7 +1,7 @@
import asyncio
+import os
import time
from collections.abc import AsyncGenerator
-import os
from typing import Callable, List, Sequence, final
import uvicorn
@@ -41,6 +41,7 @@ from exo.shared.types.tasks import ChatCompletionTaskParams
from exo.shared.types.worker.common import InstanceId
from exo.shared.types.worker.instances import Instance
+
def chunk_to_response(chunk: TokenChunk) -> ChatCompletionResponse:
return ChatCompletionResponse(
id=chunk.command_id,
@@ -49,15 +50,13 @@ def chunk_to_response(chunk: TokenChunk) -> ChatCompletionResponse:
choices=[
StreamingChoiceResponse(
index=0,
- delta=ChatCompletionMessage(
- role='assistant',
- content=chunk.text
- ),
- finish_reason=chunk.finish_reason
+ delta=ChatCompletionMessage(role="assistant", content=chunk.text),
+ finish_reason=chunk.finish_reason,
)
- ]
+ ],
)
+
async def resolve_model_meta(model_id: str) -> ModelMetadata:
if model_id in MODEL_CARDS:
model_card = MODEL_CARDS[model_id]
@@ -65,9 +64,15 @@ async def resolve_model_meta(model_id: str) -> ModelMetadata:
else:
return await get_model_meta(model_id)
+
@final
class API:
- def __init__(self, command_buffer: List[Command], global_events: AsyncSQLiteEventStorage, get_state: Callable[[], State]) -> None:
+ def __init__(
+ self,
+ command_buffer: List[Command],
+ global_events: AsyncSQLiteEventStorage,
+ get_state: Callable[[], State],
+ ) -> None:
self.get_state = get_state
self.command_buffer = command_buffer
self.global_events = global_events
@@ -76,7 +81,11 @@ class API:
self._setup_cors()
self._setup_routes()
- self._app.mount("/", StaticFiles(directory=os.environ["DASHBOARD_DIR"], html=True), name="dashboard")
+ self._app.mount(
+ "/",
+ StaticFiles(directory=os.environ["DASHBOARD_DIR"], html=True),
+ name="dashboard",
+ )
def _setup_cors(self) -> None:
self._app.add_middleware(
@@ -100,15 +109,17 @@ class API:
def app(self) -> FastAPI:
return self._app
- async def create_instance(self, payload: CreateInstanceTaskParams) -> CreateInstanceResponse:
+ async def create_instance(
+ self, payload: CreateInstanceTaskParams
+ ) -> CreateInstanceResponse:
model_meta = await resolve_model_meta(payload.model_id)
required_memory_bytes = model_meta.storage_size_kilobytes * 1024
available_memory_bytes = self._calculate_total_available_memory()
-
+
if required_memory_bytes > available_memory_bytes:
raise HTTPException(
- status_code=400,
- detail=f"Insufficient memory to create instance. Required: {required_memory_bytes // (1024**3):.1f}GB, Available: {available_memory_bytes // (1024**3):.1f}GB"
+ status_code=400,
+ detail=f"Insufficient memory to create instance. Required: {required_memory_bytes // (1024**3):.1f}GB, Available: {available_memory_bytes // (1024**3):.1f}GB",
)
command = CreateInstanceCommand(
@@ -148,7 +159,9 @@ class API:
instance_id=instance_id,
)
- async def _generate_chat_stream(self, command_id: CommandId) -> AsyncGenerator[str, None]:
+ async def _generate_chat_stream(
+ self, command_id: CommandId
+ ) -> AsyncGenerator[str, None]:
"""Generate chat completion stream as JSON strings."""
events = await self.global_events.get_events_since(0)
@@ -158,7 +171,9 @@ class API:
while not finished:
await asyncio.sleep(0.01)
- events: Sequence[EventFromEventLog[Event]] = await self.global_events.get_events_since(prev_idx)
+ events: Sequence[
+ EventFromEventLog[Event]
+ ] = await self.global_events.get_events_since(prev_idx)
# TODO: Can do this with some better functionality to tail event log into an AsyncGenerator.
prev_idx = events[-1].idx_in_log if events else prev_idx
@@ -166,17 +181,17 @@ class API:
event = wrapped_event.event
if isinstance(event, ChunkGenerated) and event.command_id == command_id:
assert isinstance(event.chunk, TokenChunk)
- chunk_response: ChatCompletionResponse = chunk_to_response(event.chunk)
+ chunk_response: ChatCompletionResponse = chunk_to_response(
+ event.chunk
+ )
print(chunk_response)
yield f"data: {chunk_response.model_dump_json()}\n\n"
if event.chunk.finish_reason is not None:
yield "data: [DONE]"
finished = True
-
- command = TaskFinishedCommand(
- command_id=command_id
- )
+
+ command = TaskFinishedCommand(command_id=command_id)
self.command_buffer.append(command)
return
@@ -184,23 +199,30 @@ class API:
async def _trigger_notify_user_to_download_model(self, model_id: str) -> None:
print("TODO: we should send a notification to the user to download the model")
- async def chat_completions(self, payload: ChatCompletionTaskParams) -> StreamingResponse:
+ async def chat_completions(
+ self, payload: ChatCompletionTaskParams
+ ) -> StreamingResponse:
"""Handle chat completions with proper streaming response."""
model_meta = await resolve_model_meta(payload.model)
payload.model = model_meta.model_id
-
+
# Preprocess messages for GPT-OSS harmony format if needed
if "gpt-oss" in payload.model.lower():
import re
+
for message in payload.messages:
if message.content and "<|channel|>" in message.content:
# Parse harmony format tags
- thinking_pattern = r'<\|channel\|>(.*?)(?=<\|message\|>|$)'
- content_pattern = r'<\|message\|>(.*?)(?=<\|end\|>|$)'
-
- thinking_match = re.search(thinking_pattern, message.content, re.DOTALL)
- content_match = re.search(content_pattern, message.content, re.DOTALL)
-
+ thinking_pattern = r"<\|channel\|>(.*?)(?=<\|message\|>|$)"
+ content_pattern = r"<\|message\|>(.*?)(?=<\|end\|>|$)"
+
+ thinking_match = re.search(
+ thinking_pattern, message.content, re.DOTALL
+ )
+ content_match = re.search(
+ content_pattern, message.content, re.DOTALL
+ )
+
if content_match:
# Extract the actual content
message.content = content_match.group(1).strip()
@@ -213,7 +235,9 @@ class API:
break
else:
await self._trigger_notify_user_to_download_model(payload.model)
- raise HTTPException(status_code=404, detail=f"No instance found for model {payload.model}")
+ raise HTTPException(
+ status_code=404, detail=f"No instance found for model {payload.model}"
+ )
command = ChatCompletionCommand(
command_id=CommandId(),
@@ -222,29 +246,33 @@ class API:
)
self.command_buffer.append(command)
return StreamingResponse(
- self._generate_chat_stream(command.command_id),
- media_type="text/plain"
+ self._generate_chat_stream(command.command_id), media_type="text/plain"
)
def _calculate_total_available_memory(self) -> int:
"""Calculate total available memory across all nodes in bytes."""
state = self.get_state()
total_available = 0
-
+
for node_profile in state.node_profiles.values():
total_available += node_profile.memory.ram_available
-
+
return total_available
async def get_models(self) -> ModelList:
"""Returns list of available models."""
- return ModelList(data=[
- ModelListModel(
- id=card.short_id,
- hugging_face_id=card.model_id,
- name=card.name,
- description=card.description,
- tags=card.tags) for card in MODEL_CARDS.values()])
+ return ModelList(
+ data=[
+ ModelListModel(
+ id=card.short_id,
+ hugging_face_id=card.model_id,
+ name=card.name,
+ description=card.description,
+ tags=card.tags,
+ )
+ for card in MODEL_CARDS.values()
+ ]
+ )
def start_fastapi_server(
diff --git a/src/exo/master/election_callback.py b/src/exo/master/election_callback.py
index 61e7c7e6..92569f3b 100644
--- a/src/exo/master/election_callback.py
+++ b/src/exo/master/election_callback.py
@@ -8,17 +8,17 @@ class ElectionCallbacks:
Simple callbacks for the Rust election system to invoke.
No event system involvement - just direct forwarder control.
"""
-
+
def __init__(self, forwarder_supervisor: ForwarderSupervisor, logger: Logger):
self._forwarder_supervisor = forwarder_supervisor
self._logger = logger
-
+
async def on_became_master(self) -> None:
"""Called when this node is elected as master"""
self._logger.info("Node elected as master")
await self._forwarder_supervisor.notify_role_change(ForwarderRole.MASTER)
-
+
async def on_became_replica(self) -> None:
"""Called when this node becomes a replica"""
self._logger.info("Node demoted to replica")
- await self._forwarder_supervisor.notify_role_change(ForwarderRole.REPLICA)
\ No newline at end of file
+ await self._forwarder_supervisor.notify_role_change(ForwarderRole.REPLICA)
diff --git a/src/exo/master/forwarder_supervisor.py b/src/exo/master/forwarder_supervisor.py
index a8f5bba1..50798c8a 100644
--- a/src/exo/master/forwarder_supervisor.py
+++ b/src/exo/master/forwarder_supervisor.py
@@ -16,6 +16,7 @@ from exo.shared.types.common import NodeId
class ForwarderRole(str, Enum):
"""Role determines which forwarding pairs to use"""
+
MASTER = "master"
REPLICA = "replica"
@@ -24,23 +25,23 @@ class ForwarderSupervisor:
"""
Manages the forwarder subprocess for SQLite ↔ libp2p event forwarding.
The forwarder is a single process that handles multiple forwarding pairs.
-
+
Master mode forwards:
- sqlite:worker_events.db:events → libp2p:worker_events (share local worker events)
- - libp2p:worker_events → sqlite:global_events.db:events (collect network worker events)
+ - libp2p:worker_events → sqlite:global_events.db:events (collect network worker events)
- sqlite:global_events.db:events → libp2p:global_events (broadcast merged global log)
-
+
Replica mode forwards:
- sqlite:worker_events.db:events → libp2p:worker_events (share local worker events)
- libp2p:global_events → sqlite:global_events.db:events (receive global log from master)
"""
-
+
def __init__(
- self,
+ self,
node_id: NodeId,
forwarder_binary_path: Path,
logger: Logger,
- health_check_interval: float = 5.0
+ health_check_interval: float = 5.0,
):
self.node_id = node_id
self._binary_path = forwarder_binary_path
@@ -49,7 +50,7 @@ class ForwarderSupervisor:
self._current_role: ForwarderRole | None = None
self._process: asyncio.subprocess.Process | None = None
self._health_check_task: asyncio.Task[None] | None = None
-
+
async def notify_role_change(self, new_role: ForwarderRole) -> None:
"""
Called by external systems (e.g., election handler) when role changes.
@@ -58,32 +59,32 @@ class ForwarderSupervisor:
if self._current_role == new_role:
self._logger.debug(f"Role unchanged: {new_role}")
return
-
+
self._logger.info(f"Role changing from {self._current_role} to {new_role}")
self._current_role = new_role
await self._restart_with_role(new_role)
-
+
async def start_as_replica(self) -> None:
"""Convenience method to start in replica mode"""
await self.notify_role_change(ForwarderRole.REPLICA)
-
+
async def stop(self) -> None:
"""Stop forwarder and cleanup"""
await self._stop_process()
self._current_role = None
-
+
def _get_forwarding_pairs(self, role: ForwarderRole) -> str:
"""
Generate forwarding pairs based on role.
Returns list of "source,sink" strings.
"""
pairs: list[str] = []
-
+
# Both master and replica forward local worker events to network
pairs.append(
f"sqlite:{EXO_WORKER_EVENT_DB}:events|libp2p:{LIBP2P_WORKER_EVENTS_TOPIC}"
)
-
+
if role == ForwarderRole.MASTER:
# Master: collect worker events from network into global log
pairs.append(
@@ -98,33 +99,31 @@ class ForwarderSupervisor:
pairs.append(
f"libp2p:{LIBP2P_GLOBAL_EVENTS_TOPIC}|sqlite:{EXO_GLOBAL_EVENT_DB}:events"
)
-
- return ','.join(pairs)
-
+
+ return ",".join(pairs)
+
async def _restart_with_role(self, role: ForwarderRole) -> None:
"""Internal method to restart forwarder with new role"""
await self._stop_process()
-
-
+
pairs: str = self._get_forwarding_pairs(role)
env_vars = os.environ.copy()
env_vars["FORWARDER_NODE_ID"] = str(self.node_id)
self._process = await asyncio.create_subprocess_exec(
str(self._binary_path),
- "--events-db", str(EXO_WORKER_EVENT_DB),
- f'{pairs}',
+ "--events-db",
+ str(EXO_WORKER_EVENT_DB),
+ f"{pairs}",
stdout=None,
stderr=None,
- env=env_vars
+ env=env_vars,
)
-
+
self._logger.info(f"Starting forwarder with forwarding pairs: {pairs}")
-
+
# Start health monitoring
- self._health_check_task = asyncio.create_task(
- self._monitor_health()
- )
-
+ self._health_check_task = asyncio.create_task(self._monitor_health())
+
async def _stop_process(self) -> None:
"""Stop the forwarder process gracefully"""
if self._health_check_task:
@@ -132,7 +131,7 @@ class ForwarderSupervisor:
with contextlib.suppress(asyncio.CancelledError):
await self._health_check_task
self._health_check_task = None
-
+
if self._process:
# Check if process is already dead
if self._process.returncode is None:
@@ -148,46 +147,45 @@ class ForwarderSupervisor:
# Process already dead
pass
self._process = None
-
+
async def _monitor_health(self) -> None:
"""Monitor process health and restart if it crashes"""
while self._process and self._current_role:
try:
# Check if process is still alive
retcode = await asyncio.wait_for(
- self._process.wait(),
- timeout=self._health_check_interval
+ self._process.wait(), timeout=self._health_check_interval
)
# Process exited
self._logger.error(f"Forwarder exited with code {retcode}")
-
+
# Auto-restart
await asyncio.sleep(0.2) # Brief delay before restart
if self._current_role: # Still have a role
await self._restart_with_role(self._current_role)
break
-
+
except asyncio.TimeoutError:
# Process still running, continue monitoring
continue
except asyncio.CancelledError:
break
-
+
@property
def is_running(self) -> bool:
"""Check if forwarder process is running"""
return self._process is not None and self._process.returncode is None
-
+
@property
def current_role(self) -> ForwarderRole | None:
"""Get current forwarder role (for testing)"""
return self._current_role
-
+
@property
def process_pid(self) -> int | None:
"""Get current process PID (for testing)"""
return self._process.pid if self._process else None
-
+
@property
def process(self) -> asyncio.subprocess.Process | None:
"""Get current process (for testing)"""
diff --git a/src/exo/master/main.py b/src/exo/master/main.py
index 6c1fc038..c0709db2 100644
--- a/src/exo/master/main.py
+++ b/src/exo/master/main.py
@@ -38,9 +38,16 @@ from exo.shared.utils import Keypair, get_node_id_keypair
class Master:
- def __init__(self, node_id_keypair: Keypair, node_id: NodeId, command_buffer: list[Command],
- global_events: AsyncSQLiteEventStorage, worker_events: AsyncSQLiteEventStorage,
- forwarder_binary_path: Path, logger: logging.Logger):
+ def __init__(
+ self,
+ node_id_keypair: Keypair,
+ node_id: NodeId,
+ command_buffer: list[Command],
+ global_events: AsyncSQLiteEventStorage,
+ worker_events: AsyncSQLiteEventStorage,
+ forwarder_binary_path: Path,
+ logger: logging.Logger,
+ ):
self.state = State()
self.node_id_keypair = node_id_keypair
self.node_id = node_id
@@ -49,9 +56,7 @@ class Master:
self.worker_events = worker_events
self.command_task_mapping: dict[CommandId, TaskId] = {}
self.forwarder_supervisor = ForwarderSupervisor(
- self.node_id,
- forwarder_binary_path=forwarder_binary_path,
- logger=logger
+ self.node_id, forwarder_binary_path=forwarder_binary_path, logger=logger
)
self.election_callbacks = ElectionCallbacks(self.forwarder_supervisor, logger)
self.logger = logger
@@ -74,7 +79,10 @@ class Master:
async def _run_event_loop_body(self) -> None:
next_events: list[Event] = []
# 1. process commands
- if self.forwarder_supervisor.current_role == ForwarderRole.MASTER and len(self.command_buffer) > 0:
+ if (
+ self.forwarder_supervisor.current_role == ForwarderRole.MASTER
+ and len(self.command_buffer) > 0
+ ):
# for now we do one command at a time
next_command = self.command_buffer.pop(0)
self.logger.info(f"got command: {next_command}")
@@ -83,43 +91,64 @@ class Master:
case ChatCompletionCommand():
matching_instance: Instance | None = None
for instance in self.state.instances.values():
- if instance.shard_assignments.model_id == next_command.request_params.model:
+ if (
+ instance.shard_assignments.model_id
+ == next_command.request_params.model
+ ):
matching_instance = instance
break
if not matching_instance:
- raise ValueError(f"No instance found for model {next_command.request_params.model}")
+ raise ValueError(
+ f"No instance found for model {next_command.request_params.model}"
+ )
task_id = TaskId()
- next_events.append(TaskCreated(
- task_id=task_id,
- task=ChatCompletionTask(
- task_type=TaskType.CHAT_COMPLETION,
+ next_events.append(
+ TaskCreated(
task_id=task_id,
- command_id=next_command.command_id,
- instance_id=matching_instance.instance_id,
- task_status=TaskStatus.PENDING,
- task_params=next_command.request_params
+ task=ChatCompletionTask(
+ task_type=TaskType.CHAT_COMPLETION,
+ task_id=task_id,
+ command_id=next_command.command_id,
+ instance_id=matching_instance.instance_id,
+ task_status=TaskStatus.PENDING,
+ task_params=next_command.request_params,
+ ),
)
- ))
+ )
self.command_task_mapping[next_command.command_id] = task_id
case DeleteInstanceCommand():
- placement = get_instance_placements(next_command, self.state.topology, self.state.instances)
- transition_events = get_transition_events(self.state.instances, placement)
+ placement = get_instance_placements(
+ next_command, self.state.topology, self.state.instances
+ )
+ transition_events = get_transition_events(
+ self.state.instances, placement
+ )
next_events.extend(transition_events)
case CreateInstanceCommand():
- placement = get_instance_placements(next_command, self.state.topology, self.state.instances)
- transition_events = get_transition_events(self.state.instances, placement)
+ placement = get_instance_placements(
+ next_command, self.state.topology, self.state.instances
+ )
+ transition_events = get_transition_events(
+ self.state.instances, placement
+ )
next_events.extend(transition_events)
case TaskFinishedCommand():
- next_events.append(TaskDeleted(
- task_id=self.command_task_mapping[next_command.command_id]
- ))
+ next_events.append(
+ TaskDeleted(
+ task_id=self.command_task_mapping[next_command.command_id]
+ )
+ )
del self.command_task_mapping[next_command.command_id]
- await self.event_log_for_writes.append_events(next_events, origin=self.node_id)
+ await self.event_log_for_writes.append_events(
+ next_events, origin=self.node_id
+ )
# 2. get latest events
- events = await self.event_log_for_reads.get_events_since(self.state.last_event_applied_idx, ignore_no_op_events=True)
+ events = await self.event_log_for_reads.get_events_since(
+ self.state.last_event_applied_idx, ignore_no_op_events=True
+ )
if len(events) == 0:
await asyncio.sleep(0.01)
return
@@ -133,8 +162,15 @@ class Master:
# TODO: This can be done in a better place. But for now, we use this to check if any running instances have been broken.
write_events: list[Event] = []
- if any([isinstance(event_from_log.event, TopologyEdgeDeleted) for event_from_log in events]):
- connected_node_ids = set([x.node_id for x in self.state.topology.list_nodes()])
+ if any(
+ [
+ isinstance(event_from_log.event, TopologyEdgeDeleted)
+ for event_from_log in events
+ ]
+ ):
+ connected_node_ids = set(
+ [x.node_id for x in self.state.topology.list_nodes()]
+ )
for instance_id, instance in self.state.instances.items():
delete = False
for node_id in instance.shard_assignments.node_to_runner:
@@ -142,31 +178,40 @@ class Master:
delete = True
break
if delete:
- write_events.append(InstanceDeleted(
- instance_id=instance_id
- ))
+ write_events.append(InstanceDeleted(instance_id=instance_id))
if write_events:
- await self.event_log_for_writes.append_events(events=write_events, origin=self.node_id)
+ await self.event_log_for_writes.append_events(
+ events=write_events, origin=self.node_id
+ )
async def run(self):
self.state = await self._get_state_snapshot()
-
+
async def heartbeat_task():
while True:
- await self.event_log_for_writes.append_events([Heartbeat(node_id=self.node_id)], origin=self.node_id)
+ await self.event_log_for_writes.append_events(
+ [Heartbeat(node_id=self.node_id)], origin=self.node_id
+ )
await asyncio.sleep(5)
+
asyncio.create_task(heartbeat_task())
# TODO: we should clean these up on shutdown
await self.forwarder_supervisor.start_as_replica()
- if os.getenv('EXO_RUN_AS_REPLICA') in set(['TRUE', 'true', '1']):
+ if os.getenv("EXO_RUN_AS_REPLICA") in set(["TRUE", "true", "1"]):
await self.election_callbacks.on_became_replica()
else:
await self.election_callbacks.on_became_master()
- role = "MASTER" if self.forwarder_supervisor.current_role == ForwarderRole.MASTER else "REPLICA"
- await self.event_log_for_writes.append_events([TopologyNodeCreated(node_id=self.node_id, role=role)], origin=self.node_id)
+ role = (
+ "MASTER"
+ if self.forwarder_supervisor.current_role == ForwarderRole.MASTER
+ else "REPLICA"
+ )
+ await self.event_log_for_writes.append_events(
+ [TopologyNodeCreated(node_id=self.node_id, role=role)], origin=self.node_id
+ )
while True:
try:
await self._run_event_loop_body()
@@ -177,11 +222,13 @@ class Master:
async def async_main():
- logger = logging.getLogger('master_logger')
+ logger = logging.getLogger("master_logger")
logger.setLevel(logging.INFO)
if not logger.handlers:
handler = logging.StreamHandler()
- handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
+ handler.setFormatter(
+ logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
+ )
logger.addHandler(handler)
node_id_keypair = get_node_id_keypair()
@@ -203,19 +250,28 @@ async def async_main():
global_events,
lambda: master.state,
"0.0.0.0",
- int(os.environ.get("API_PORT", 8000))
+ int(os.environ.get("API_PORT", 8000)),
),
- daemon=True
+ daemon=True,
)
api_thread.start()
- logger.info('Running FastAPI server in a separate thread. Listening on port 8000.')
+ logger.info("Running FastAPI server in a separate thread. Listening on port 8000.")
- master = Master(node_id_keypair, node_id, command_buffer, global_events, worker_events,
- Path(os.environ["GO_BUILD_DIR"])/"forwarder", logger)
+ master = Master(
+ node_id_keypair,
+ node_id,
+ command_buffer,
+ global_events,
+ worker_events,
+ Path(os.environ["GO_BUILD_DIR"]) / "forwarder",
+ logger,
+ )
await master.run()
+
def main():
asyncio.run(async_main())
+
if __name__ == "__main__":
main()
diff --git a/src/exo/master/placement.py b/src/exo/master/placement.py
index e047cfa0..f61da749 100644
--- a/src/exo/master/placement.py
+++ b/src/exo/master/placement.py
@@ -11,9 +11,12 @@ from exo.master.utils.placement_utils import (
get_smallest_cycles,
)
from exo.shared.topology import Topology
-from exo.shared.types.common import Host
+from exo.shared.types.common import Host
from exo.shared.types.events import Event, InstanceCreated, InstanceDeleted
-from exo.shared.types.events.commands import CreateInstanceCommand, DeleteInstanceCommand
+from exo.shared.types.events.commands import (
+ CreateInstanceCommand,
+ DeleteInstanceCommand,
+)
from exo.shared.types.worker.common import InstanceId
from exo.shared.types.worker.instances import Instance, InstanceStatus
@@ -21,62 +24,85 @@ from exo.shared.types.worker.instances import Instance, InstanceStatus
def random_ephemeral_port() -> int:
return random.randint(49152, 65535)
+
@singledispatch
def get_instance_placements(
command: CreateInstanceCommand,
topology: Topology,
current_instances: dict[InstanceId, Instance],
) -> dict[InstanceId, Instance]:
- available_models = [current_instances[instance].shard_assignments.model_id for instance in current_instances]
+ available_models = [
+ current_instances[instance].shard_assignments.model_id
+ for instance in current_instances
+ ]
if command.model_meta.model_id in available_models:
raise ValueError(f"Instance for {command.model_meta.model_id} already exists")
-
+
all_nodes = list(topology.list_nodes())
cycles = topology.get_cycles()
# we can also always just have a node on its own
singleton_cycles = [[node] for node in all_nodes]
candidate_cycles = cycles + singleton_cycles
- cycles_with_sufficient_memory = filter_cycles_by_memory(candidate_cycles, command.model_meta.storage_size_kilobytes * 1024)
+ cycles_with_sufficient_memory = filter_cycles_by_memory(
+ candidate_cycles, command.model_meta.storage_size_kilobytes * 1024
+ )
if not cycles_with_sufficient_memory:
raise ValueError("No cycles found with sufficient memory")
smallest_cycles = get_smallest_cycles(cycles_with_sufficient_memory)
selected_cycle = None
- has_thunderbolt_cycle = any([
- topology.get_subgraph_from_nodes(cycle).is_thunderbolt_cycle(cycle)
- for cycle in smallest_cycles
- ])
+ has_thunderbolt_cycle = any(
+ [
+ topology.get_subgraph_from_nodes(cycle).is_thunderbolt_cycle(cycle)
+ for cycle in smallest_cycles
+ ]
+ )
if has_thunderbolt_cycle:
smallest_cycles = [
- cycle for cycle in smallest_cycles
+ cycle
+ for cycle in smallest_cycles
if topology.get_subgraph_from_nodes(cycle).is_thunderbolt_cycle(cycle)
]
- selected_cycle = max(smallest_cycles, key=lambda cycle: sum(node.node_profile.memory.ram_available for node in cycle if node.node_profile is not None))
-
+ selected_cycle = max(
+ smallest_cycles,
+ key=lambda cycle: sum(
+ node.node_profile.memory.ram_available
+ for node in cycle
+ if node.node_profile is not None
+ ),
+ )
+
shard_assignments = get_shard_assignments(command.model_meta, selected_cycle)
-
+
cycle_digraph: Topology = topology.get_subgraph_from_nodes(selected_cycle)
hosts: list[Host] = get_hosts_from_subgraph(cycle_digraph)
-
+
instance_id = command.instance_id
target_instances = deepcopy(current_instances)
target_instances[instance_id] = Instance(
instance_id=instance_id,
instance_type=InstanceStatus.ACTIVE,
shard_assignments=shard_assignments,
- hosts=[Host(
- ip=host.ip,
- # NOTE: it's fine to have non-deterministic ports here since this is in a command decision
- port=random_ephemeral_port(),
- ) for host in hosts]
+ hosts=[
+ Host(
+ ip=host.ip,
+ # NOTE: it's fine to have non-deterministic ports here since this is in a command decision
+ port=random_ephemeral_port(),
+ )
+ for host in hosts
+ ],
)
return target_instances
@get_instance_placements.register
-def _(command: DeleteInstanceCommand, topology: Topology, current_instances: dict[InstanceId, Instance]) -> dict[InstanceId, Instance]:
+def _(
+ command: DeleteInstanceCommand,
+ topology: Topology,
+ current_instances: dict[InstanceId, Instance],
+) -> dict[InstanceId, Instance]:
target_instances = deepcopy(current_instances)
if command.instance_id in target_instances:
del target_instances[command.instance_id]
@@ -107,5 +133,5 @@ def get_transition_events(
instance_id=instance_id,
)
)
-
+
return events
diff --git a/src/exo/master/tests/api_utils_test.py b/src/exo/master/tests/api_utils_test.py
index 0b3a666a..5682f0e5 100644
--- a/src/exo/master/tests/api_utils_test.py
+++ b/src/exo/master/tests/api_utils_test.py
@@ -27,8 +27,9 @@ _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]]
+ func: Callable[_P, Awaitable[_R]],
) -> Callable[_P, Coroutine[Any, Any, _R]]:
@pytest.mark.asyncio
@functools.wraps(func)
@@ -40,11 +41,14 @@ def with_master_main(
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
@@ -59,6 +63,7 @@ class ChatMessage:
else:
raise ValueError(f"Unsupported role: {self.role}")
+
async def stream_chatgpt_response(
messages: list[ChatMessage],
model: str = "gpt-3.5-turbo",
@@ -67,7 +72,9 @@ async def stream_chatgpt_response(
api_key=OPENAI_API_KEY,
base_url=OPENAI_API_URL,
)
- openai_messages: list[ChatCompletionMessageParam] = [m.to_openai() for m in messages]
+ openai_messages: list[ChatCompletionMessageParam] = [
+ m.to_openai() for m in messages
+ ]
stream: AsyncStream[ChatCompletionChunk] = await client.chat.completions.create(
model=model,
messages=openai_messages,
diff --git a/src/exo/master/tests/conftest.py b/src/exo/master/tests/conftest.py
index f951d802..fcfaace4 100644
--- a/src/exo/master/tests/conftest.py
+++ b/src/exo/master/tests/conftest.py
@@ -25,11 +25,11 @@ def create_node():
ram_total=1000,
ram_available=memory,
swap_total=1000,
- swap_available=1000
+ swap_available=1000,
),
network_interfaces=[],
- system=SystemPerformanceProfile(flops_fp16=1000)
- )
+ system=SystemPerformanceProfile(flops_fp16=1000),
+ ),
)
return _create_node
@@ -39,7 +39,10 @@ def create_node():
@pytest.fixture
def create_connection():
port_counter = 1235
- def _create_connection(source_node_id: NodeId, sink_node_id: NodeId, send_back_port: int | None = None) -> Connection:
+
+ def _create_connection(
+ source_node_id: NodeId, sink_node_id: NodeId, send_back_port: int | None = None
+ ) -> Connection:
nonlocal port_counter
if send_back_port is None:
send_back_port = port_counter
@@ -48,8 +51,12 @@ def create_connection():
local_node_id=source_node_id,
send_back_node_id=sink_node_id,
local_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/1234"),
- send_back_multiaddr=Multiaddr(address=f"/ip4/127.0.0.1/tcp/{send_back_port}"),
- connection_profile=ConnectionProfile(throughput=1000, latency=1000, jitter=1000)
+ send_back_multiaddr=Multiaddr(
+ address=f"/ip4/127.0.0.1/tcp/{send_back_port}"
+ ),
+ connection_profile=ConnectionProfile(
+ throughput=1000, latency=1000, jitter=1000
+ ),
)
return _create_connection
diff --git a/src/exo/master/tests/test_api.py b/src/exo/master/tests/test_api.py
index a0867c3a..ce9e1376 100644
--- a/src/exo/master/tests/test_api.py
+++ b/src/exo/master/tests/test_api.py
@@ -14,9 +14,7 @@ from exo.master.tests.api_utils_test import (
async def test_master_api_multiple_response_sequential() -> None:
# TODO: This hangs at the moment it seems.
return
- messages = [
- ChatMessage(role="user", content="Hello, who are you?")
- ]
+ messages = [ChatMessage(role="user", content="Hello, who are you?")]
token_count = 0
text: str = ""
async for choice in stream_chatgpt_response(messages):
@@ -30,11 +28,9 @@ async def test_master_api_multiple_response_sequential() -> None:
assert token_count >= 3, f"Expected at least 3 tokens, got {token_count}"
assert len(text) > 0, "Expected non-empty response text"
- await asyncio.sleep(0.1)
+ await asyncio.sleep(0.1)
- messages = [
- ChatMessage(role="user", content="What time is it in France?")
- ]
+ messages = [ChatMessage(role="user", content="What time is it in France?")]
token_count = 0
text = "" # re-initialize, do not redeclare type
async for choice in stream_chatgpt_response(messages):
diff --git a/src/exo/master/tests/test_forwarder_supervisor.py b/src/exo/master/tests/test_forwarder_supervisor.py
index 295f6039..00829696 100644
--- a/src/exo/master/tests/test_forwarder_supervisor.py
+++ b/src/exo/master/tests/test_forwarder_supervisor.py
@@ -2,6 +2,7 @@
Comprehensive unit tests for Forwardersupervisor.
Tests basic functionality, process management, and edge cases.
"""
+
import asyncio
import logging
import os
@@ -105,6 +106,7 @@ def temp_dir() -> Generator[Path, None, None]:
yield temp_path
# Clean up
import shutil
+
shutil.rmtree(temp_path, ignore_errors=True)
@@ -122,15 +124,17 @@ def test_logger() -> logging.Logger:
"""Create a test logger."""
logger = logging.getLogger("test_forwarder")
logger.setLevel(logging.DEBUG)
-
+
# Add console handler for debugging
if not logger.handlers:
handler = logging.StreamHandler()
handler.setLevel(logging.DEBUG)
- formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+ formatter = logging.Formatter(
+ "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+ )
handler.setFormatter(formatter)
logger.addHandler(handler)
-
+
return logger
@@ -147,69 +151,76 @@ def mock_env_vars(temp_dir: Path) -> dict[str, str]:
async def cleanup_processes() -> AsyncGenerator[set[int], None]:
"""Track and cleanup any processes created during tests."""
tracked_pids: set[int] = set()
-
+
yield tracked_pids
-
+
# Cleanup any remaining processes - simplified to avoid psutil dependency
import contextlib
import subprocess
+
for pid in tracked_pids:
with contextlib.suppress(Exception):
subprocess.run(["kill", str(pid)], check=False, timeout=1)
@pytest.fixture
-def track_subprocess(cleanup_processes: set[int]) -> Callable[[asyncio.subprocess.Process], asyncio.subprocess.Process]:
+def track_subprocess(
+ cleanup_processes: set[int],
+) -> Callable[[asyncio.subprocess.Process], asyncio.subprocess.Process]:
"""Function to track subprocess PIDs for cleanup."""
+
def track(process: asyncio.subprocess.Process) -> asyncio.subprocess.Process:
if process.pid:
cleanup_processes.add(process.pid)
return process
+
return track
class TestForwardersupervisorBasic:
"""Basic functionality tests for Forwardersupervisor."""
-
+
@pytest.mark.asyncio
async def test_start_as_replica(
self,
mock_forwarder_script: Path,
mock_env_vars: dict[str, str],
test_logger: logging.Logger,
- track_subprocess: Callable[[asyncio.subprocess.Process], asyncio.subprocess.Process]
+ track_subprocess: Callable[
+ [asyncio.subprocess.Process], asyncio.subprocess.Process
+ ],
) -> None:
"""Test starting forwarder in replica mode."""
# Set environment
os.environ.update(mock_env_vars)
-
+
supervisor = ForwarderSupervisor(NodeId(), mock_forwarder_script, test_logger)
await supervisor.start_as_replica()
-
+
# Track the process for cleanup
if supervisor.process:
track_subprocess(supervisor.process)
-
+
try:
# Verify process is running
assert supervisor.is_running
assert supervisor.current_role == ForwarderRole.REPLICA
-
+
# Wait a bit for log file to be written
await asyncio.sleep(0.5)
-
+
# Verify forwarding pairs in log
log_content = Path(mock_env_vars["MOCK_LOG_FILE"]).read_text()
-
+
# Expected replica forwarding pairs
expected_pairs = [
f"sqlite:{EXO_WORKER_EVENT_DB}:events|libp2p:{LIBP2P_WORKER_EVENTS_TOPIC}",
- f"libp2p:{LIBP2P_GLOBAL_EVENTS_TOPIC}|sqlite:{EXO_GLOBAL_EVENT_DB}:events"
+ f"libp2p:{LIBP2P_GLOBAL_EVENTS_TOPIC}|sqlite:{EXO_GLOBAL_EVENT_DB}:events",
]
-
+
# Check that the forwarder received the correct arguments
assert all(pair in log_content for pair in expected_pairs)
-
+
finally:
await supervisor.stop()
assert not supervisor.is_running
@@ -220,41 +231,43 @@ class TestForwardersupervisorBasic:
mock_forwarder_script: Path,
mock_env_vars: dict[str, str],
test_logger: logging.Logger,
- track_subprocess: Callable[[asyncio.subprocess.Process], asyncio.subprocess.Process]
+ track_subprocess: Callable[
+ [asyncio.subprocess.Process], asyncio.subprocess.Process
+ ],
) -> None:
"""Test changing role from replica to master."""
os.environ.update(mock_env_vars)
-
+
supervisor = ForwarderSupervisor(NodeId(), mock_forwarder_script, test_logger)
await supervisor.start_as_replica()
-
+
if supervisor.process:
track_subprocess(supervisor.process)
-
+
try:
# Change to master
await supervisor.notify_role_change(ForwarderRole.MASTER)
-
+
if supervisor.process:
track_subprocess(supervisor.process)
-
+
# Wait for restart
await asyncio.sleep(0.5)
-
+
assert supervisor.is_running
assert supervisor.current_role == ForwarderRole.MASTER
-
+
# Verify new forwarding pairs
log_content = Path(mock_env_vars["MOCK_LOG_FILE"]).read_text()
-
+
# Expected master forwarding pairs
master_pairs = [
f"libp2p:{LIBP2P_WORKER_EVENTS_TOPIC}|sqlite:{EXO_GLOBAL_EVENT_DB}:events",
- f"sqlite:{EXO_GLOBAL_EVENT_DB}:events|libp2p:{LIBP2P_GLOBAL_EVENTS_TOPIC}"
+ f"sqlite:{EXO_GLOBAL_EVENT_DB}:events|libp2p:{LIBP2P_GLOBAL_EVENTS_TOPIC}",
]
-
+
assert all(pair in log_content for pair in master_pairs)
-
+
finally:
await supervisor.stop()
@@ -264,25 +277,27 @@ class TestForwardersupervisorBasic:
mock_forwarder_script: Path,
mock_env_vars: dict[str, str],
test_logger: logging.Logger,
- track_subprocess: Callable[[asyncio.subprocess.Process], asyncio.subprocess.Process],
+ track_subprocess: Callable[
+ [asyncio.subprocess.Process], asyncio.subprocess.Process
+ ],
) -> None:
"""Test that setting the same role twice doesn't restart the process."""
os.environ.update(mock_env_vars)
-
+
supervisor = ForwarderSupervisor(NodeId(), mock_forwarder_script, test_logger)
await supervisor.start_as_replica()
-
+
original_pid = supervisor.process_pid
if supervisor.process:
track_subprocess(supervisor.process)
-
+
try:
# Try to change to the same role
await supervisor.notify_role_change(ForwarderRole.REPLICA)
-
+
# Should not restart (same PID)
assert supervisor.process_pid == original_pid
-
+
finally:
await supervisor.stop()
@@ -292,64 +307,64 @@ class TestForwardersupervisorBasic:
mock_forwarder_script: Path,
mock_env_vars: dict[str, str],
test_logger: logging.Logger,
- track_subprocess: Callable[[asyncio.subprocess.Process], asyncio.subprocess.Process]
+ track_subprocess: Callable[
+ [asyncio.subprocess.Process], asyncio.subprocess.Process
+ ],
) -> None:
"""Test that Forwardersupervisor restarts the process if it crashes."""
# Configure mock to exit after 1 second
mock_env_vars["MOCK_EXIT_AFTER"] = "1"
mock_env_vars["MOCK_EXIT_CODE"] = "1"
os.environ.update(mock_env_vars)
-
+
supervisor = ForwarderSupervisor(
NodeId(),
mock_forwarder_script,
test_logger,
- health_check_interval=0.5 # Faster health checks for testing
+ health_check_interval=0.5, # Faster health checks for testing
)
await supervisor.start_as_replica()
-
+
original_pid = supervisor.process_pid
if supervisor.process:
track_subprocess(supervisor.process)
-
+
try:
# Wait for first crash
await asyncio.sleep(1.5)
-
+
# Process should have crashed
assert not supervisor.is_running or supervisor.process_pid != original_pid
-
+
# Clear the crash-inducing environment variables so restart works
if "MOCK_EXIT_AFTER" in os.environ:
del os.environ["MOCK_EXIT_AFTER"]
if "MOCK_EXIT_CODE" in os.environ:
del os.environ["MOCK_EXIT_CODE"]
-
+
# Wait for restart
await asyncio.sleep(1.0)
-
+
# Process should have restarted with new PID
assert supervisor.is_running
assert supervisor.process_pid != original_pid
-
+
# Track new process
if supervisor.process:
track_subprocess(supervisor.process)
-
+
finally:
await supervisor.stop()
@pytest.mark.asyncio
async def test_nonexistent_binary(
- self,
- test_logger: logging.Logger,
- temp_dir: Path
+ self, test_logger: logging.Logger, temp_dir: Path
) -> None:
"""Test behavior when forwarder binary doesn't exist."""
nonexistent_path = temp_dir / "nonexistent_forwarder"
-
+
supervisor = ForwarderSupervisor(NodeId(), nonexistent_path, test_logger)
-
+
# Should raise FileNotFoundError
with pytest.raises(FileNotFoundError):
await supervisor.start_as_replica()
@@ -357,16 +372,16 @@ class TestForwardersupervisorBasic:
class TestElectionCallbacks:
"""Test suite for ElectionCallbacks."""
-
+
@pytest.mark.asyncio
async def test_on_became_master(self, test_logger: logging.Logger) -> None:
"""Test callback when becoming master."""
mock_supervisor = MagicMock(spec=ForwarderSupervisor)
mock_supervisor.notify_role_change = AsyncMock()
-
+
callbacks = ElectionCallbacks(mock_supervisor, test_logger)
await callbacks.on_became_master()
-
+
mock_supervisor.notify_role_change.assert_called_once_with(ForwarderRole.MASTER) # type: ignore
@pytest.mark.asyncio
@@ -374,8 +389,10 @@ class TestElectionCallbacks:
"""Test callback when becoming replica."""
mock_supervisor = MagicMock(spec=ForwarderSupervisor)
mock_supervisor.notify_role_change = AsyncMock()
-
+
callbacks = ElectionCallbacks(mock_supervisor, test_logger)
await callbacks.on_became_replica()
-
- mock_supervisor.notify_role_change.assert_called_once_with(ForwarderRole.REPLICA) # type: ignore
\ No newline at end of file
+
+ mock_supervisor.notify_role_change.assert_called_once_with(
+ ForwarderRole.REPLICA
+ ) # type: ignore
diff --git a/src/exo/master/tests/test_master.py b/src/exo/master/tests/test_master.py
index fa32c7f3..293e454d 100644
--- a/src/exo/master/tests/test_master.py
+++ b/src/exo/master/tests/test_master.py
@@ -45,9 +45,10 @@ def _create_forwarder_dummy_binary() -> Path:
path.chmod(0o755)
return path
+
@pytest.mark.asyncio
async def test_master():
- logger = Logger(name='test_master_logger')
+ logger = Logger(name="test_master_logger")
event_log_manager = EventLogManager(EventLogConfig(), logger=logger)
await event_log_manager.initialize()
global_events: AsyncSQLiteEventStorage = event_log_manager.global_events
@@ -60,11 +61,11 @@ async def test_master():
for e in orig_events:
if isinstance(e.event, Heartbeat):
continue
- events.append(EventFromEventLog(
- event=e.event,
- origin=e.origin,
- idx_in_log=override_idx_in_log
- ))
+ events.append(
+ EventFromEventLog(
+ event=e.event, origin=e.origin, idx_in_log=override_idx_in_log
+ )
+ )
override_idx_in_log += 1
return events
@@ -74,40 +75,57 @@ async def test_master():
node_id_keypair = Keypair.generate_ed25519()
node_id = NodeId(node_id_keypair.to_peer_id().to_base58())
- master = Master(node_id_keypair, node_id, command_buffer=command_buffer, global_events=global_events,
- forwarder_binary_path=forwarder_binary_path, logger=logger, worker_events=global_events)
+ master = Master(
+ node_id_keypair,
+ node_id,
+ command_buffer=command_buffer,
+ global_events=global_events,
+ forwarder_binary_path=forwarder_binary_path,
+ logger=logger,
+ worker_events=global_events,
+ )
asyncio.create_task(master.run())
# wait for initial topology event
while len(list(master.state.topology.list_nodes())) == 0:
print("waiting")
await asyncio.sleep(0.001)
# inject a NodePerformanceProfile event
- await event_log_manager.global_events.append_events([
- NodePerformanceMeasured(
- node_id=node_id,
- node_profile=NodePerformanceProfile(
- model_id="maccy",
- chip_id="arm",
- friendly_name="test",
- memory=MemoryPerformanceProfile(ram_total=678948*1024, ram_available=678948*1024, swap_total=0, swap_available=0),
- network_interfaces=[],
- system=SystemPerformanceProfile(flops_fp16=0)
+ await event_log_manager.global_events.append_events(
+ [
+ NodePerformanceMeasured(
+ node_id=node_id,
+ node_profile=NodePerformanceProfile(
+ model_id="maccy",
+ chip_id="arm",
+ friendly_name="test",
+ memory=MemoryPerformanceProfile(
+ ram_total=678948 * 1024,
+ ram_available=678948 * 1024,
+ swap_total=0,
+ swap_available=0,
+ ),
+ network_interfaces=[],
+ system=SystemPerformanceProfile(flops_fp16=0),
+ ),
)
- )
- ], origin=node_id)
+ ],
+ origin=node_id,
+ )
while len(master.state.node_profiles) == 0:
await asyncio.sleep(0.001)
- command_buffer.append(CreateInstanceCommand(
- command_id=CommandId(),
- instance_id=InstanceId(),
- model_meta=ModelMetadata(
- model_id="llama-3.2-1b",
- pretty_name="Llama 3.2 1B",
- n_layers=16,
- storage_size_kilobytes=678948
+ command_buffer.append(
+ CreateInstanceCommand(
+ command_id=CommandId(),
+ instance_id=InstanceId(),
+ model_meta=ModelMetadata(
+ model_id="llama-3.2-1b",
+ pretty_name="Llama 3.2 1B",
+ n_layers=16,
+ storage_size_kilobytes=678948,
+ ),
)
- ))
+ )
while len(master.state.instances.keys()) == 0:
await asyncio.sleep(0.001)
command_buffer.append(
@@ -115,8 +133,10 @@ async def test_master():
command_id=CommandId(),
request_params=ChatCompletionTaskParams(
model="llama-3.2-1b",
- messages=[ChatCompletionMessage(role="user", content="Hello, how are you?")]
- )
+ messages=[
+ ChatCompletionMessage(role="user", content="Hello, how are you?")
+ ],
+ ),
)
)
while len(await _get_events()) < 4:
@@ -129,7 +149,9 @@ async def test_master():
assert isinstance(events[0].event, TopologyNodeCreated)
assert isinstance(events[1].event, NodePerformanceMeasured)
assert isinstance(events[2].event, InstanceCreated)
- runner_id = list(events[2].event.instance.shard_assignments.runner_to_shard.keys())[0]
+ runner_id = list(events[2].event.instance.shard_assignments.runner_to_shard.keys())[
+ 0
+ ]
assert events[2].event == InstanceCreated(
instance=Instance(
instance_id=events[2].event.instance.instance_id,
@@ -146,15 +168,15 @@ async def test_master():
model_id="llama-3.2-1b",
pretty_name="Llama 3.2 1B",
n_layers=16,
- storage_size_kilobytes=678948
+ storage_size_kilobytes=678948,
),
device_rank=0,
- world_size=1
+ world_size=1,
)
},
- node_to_runner={node_id: runner_id}
+ node_to_runner={node_id: runner_id},
),
- hosts=[]
+ hosts=[],
)
)
assert isinstance(events[3].event, TaskCreated)
@@ -168,8 +190,10 @@ async def test_master():
task_status=TaskStatus.PENDING,
task_params=ChatCompletionTaskParams(
model="llama-3.2-1b",
- messages=[ChatCompletionMessage(role="user", content="Hello, how are you?")]
- )
- )
+ messages=[
+ ChatCompletionMessage(role="user", content="Hello, how are you?")
+ ],
+ ),
+ ),
)
assert len(command_buffer) == 0
diff --git a/src/exo/master/tests/test_placement.py b/src/exo/master/tests/test_placement.py
index c901498d..4f83fcfa 100644
--- a/src/exo/master/tests/test_placement.py
+++ b/src/exo/master/tests/test_placement.py
@@ -20,28 +20,29 @@ from exo.shared.types.worker.runners import ShardAssignments
def topology() -> Topology:
return Topology()
+
@pytest.fixture
def instance() -> Instance:
return Instance(
instance_id=InstanceId(),
instance_type=InstanceStatus.ACTIVE,
shard_assignments=ShardAssignments(
- model_id="test-model",
- runner_to_shard={},
- node_to_runner={}
+ model_id="test-model", runner_to_shard={}, node_to_runner={}
),
- hosts=[]
+ hosts=[],
)
+
@pytest.fixture
def model_meta() -> ModelMetadata:
return ModelMetadata(
model_id="test-model",
storage_size_kilobytes=1000,
pretty_name="Test Model",
- n_layers=10
+ n_layers=10,
)
+
def create_instance_command(model_meta: ModelMetadata) -> CreateInstanceCommand:
return CreateInstanceCommand(
command_id=CommandId(),
@@ -50,11 +51,14 @@ def create_instance_command(model_meta: ModelMetadata) -> CreateInstanceCommand:
)
-@pytest.mark.parametrize("available_memory,total_layers,expected_layers", [
- ((500, 500, 1000), 12, (3, 3, 6)),
- ((500, 500, 500), 12, (4, 4, 4)),
- ((312, 518, 1024), 12, (2, 3, 7))
-])
+@pytest.mark.parametrize(
+ "available_memory,total_layers,expected_layers",
+ [
+ ((500, 500, 1000), 12, (3, 3, 6)),
+ ((500, 500, 500), 12, (4, 4, 4)),
+ ((312, 518, 1024), 12, (2, 3, 7)),
+ ],
+)
def test_get_instance_placements_create_instance(
available_memory: tuple[int, int, int],
total_layers: int,
@@ -62,12 +66,14 @@ def test_get_instance_placements_create_instance(
topology: Topology,
model_meta: ModelMetadata,
create_node: Callable[[int, NodeId | None], Node],
- create_connection: Callable[[NodeId, NodeId], Connection]
+ create_connection: Callable[[NodeId, NodeId], Connection],
):
# arrange
model_meta.n_layers = total_layers
- model_meta.storage_size_kilobytes = sum(available_memory) # make it exactly fit across all nodes
-
+ model_meta.storage_size_kilobytes = sum(
+ available_memory
+ ) # make it exactly fit across all nodes
+
create_instance_command = CreateInstanceCommand(
command_id=CommandId(),
model_meta=model_meta,
@@ -76,9 +82,9 @@ def test_get_instance_placements_create_instance(
node_id_a = NodeId()
node_id_b = NodeId()
node_id_c = NodeId()
- topology.add_node(create_node(available_memory[0]*1024, node_id_a))
- topology.add_node(create_node(available_memory[1]*1024, node_id_b))
- topology.add_node(create_node(available_memory[2]*1024, node_id_c))
+ topology.add_node(create_node(available_memory[0] * 1024, node_id_a))
+ topology.add_node(create_node(available_memory[1] * 1024, node_id_b))
+ topology.add_node(create_node(available_memory[2] * 1024, node_id_c))
topology.add_connection(create_connection(node_id_a, node_id_b))
topology.add_connection(create_connection(node_id_b, node_id_c))
topology.add_connection(create_connection(node_id_c, node_id_a))
@@ -95,33 +101,34 @@ def test_get_instance_placements_create_instance(
runner_id_a = instance.shard_assignments.node_to_runner[node_id_a]
runner_id_b = instance.shard_assignments.node_to_runner[node_id_b]
runner_id_c = instance.shard_assignments.node_to_runner[node_id_c]
-
+
shard_a = instance.shard_assignments.runner_to_shard[runner_id_a]
shard_b = instance.shard_assignments.runner_to_shard[runner_id_b]
shard_c = instance.shard_assignments.runner_to_shard[runner_id_c]
-
+
assert shard_a.end_layer - shard_a.start_layer == expected_layers[0]
assert shard_b.end_layer - shard_b.start_layer == expected_layers[1]
assert shard_c.end_layer - shard_c.start_layer == expected_layers[2]
-
+
shards = [shard_a, shard_b, shard_c]
shards_sorted = sorted(shards, key=lambda s: s.start_layer)
assert shards_sorted[0].start_layer == 0
assert shards_sorted[-1].end_layer == total_layers
+
def test_get_instance_placements_one_node_exact_fit(
create_node: Callable[[int, NodeId | None], Node],
) -> None:
topology = Topology()
node_id = NodeId()
- topology.add_node(create_node(1000*1024, node_id))
+ topology.add_node(create_node(1000 * 1024, node_id))
create_instance_command = CreateInstanceCommand(
command_id=CommandId(),
model_meta=ModelMetadata(
model_id="test-model",
storage_size_kilobytes=1000,
pretty_name="Test Model",
- n_layers=10
+ n_layers=10,
),
instance_id=InstanceId(),
)
@@ -135,19 +142,20 @@ def test_get_instance_placements_one_node_exact_fit(
assert len(instance.shard_assignments.runner_to_shard) == 1
assert len(instance.shard_assignments.runner_to_shard) == 1
+
def test_get_instance_placements_one_node_fits_with_extra_memory(
create_node: Callable[[int, NodeId | None], Node],
) -> None:
topology = Topology()
node_id = NodeId()
- topology.add_node(create_node(1001*1024, node_id))
+ topology.add_node(create_node(1001 * 1024, node_id))
create_instance_command = CreateInstanceCommand(
command_id=CommandId(),
model_meta=ModelMetadata(
model_id="test-model",
storage_size_kilobytes=1000,
pretty_name="Test Model",
- n_layers=10
+ n_layers=10,
),
instance_id=InstanceId(),
)
@@ -161,19 +169,20 @@ def test_get_instance_placements_one_node_fits_with_extra_memory(
assert len(instance.shard_assignments.runner_to_shard) == 1
assert len(instance.shard_assignments.runner_to_shard) == 1
+
def test_get_instance_placements_one_node_not_fit(
create_node: Callable[[int, NodeId | None], Node],
) -> None:
topology = Topology()
node_id = NodeId()
- topology.add_node(create_node(1000*1024, node_id))
+ topology.add_node(create_node(1000 * 1024, node_id))
create_instance_command = CreateInstanceCommand(
command_id=CommandId(),
model_meta=ModelMetadata(
model_id="test-model",
storage_size_kilobytes=1001,
pretty_name="Test Model",
- n_layers=10
+ n_layers=10,
),
instance_id=InstanceId(),
)
@@ -181,15 +190,12 @@ def test_get_instance_placements_one_node_not_fit(
with pytest.raises(ValueError, match="No cycles found with sufficient memory"):
get_instance_placements(create_instance_command, topology, {})
+
def test_get_transition_events_no_change(topology: Topology, instance: Instance):
# arrange
instance_id = InstanceId()
- current_instances = {
- instance_id: instance
- }
- target_instances = {
- instance_id: instance
- }
+ current_instances = {instance_id: instance}
+ target_instances = {instance_id: instance}
# act
events = get_transition_events(current_instances, target_instances)
@@ -202,9 +208,7 @@ def test_get_transition_events_create_instance(topology: Topology, instance: Ins
# arrange
instance_id = InstanceId()
current_instances: dict[InstanceId, Instance] = {}
- target_instances: dict[InstanceId, Instance] = {
- instance_id: instance
- }
+ target_instances: dict[InstanceId, Instance] = {instance_id: instance}
# act
events = get_transition_events(current_instances, target_instances)
@@ -217,9 +221,7 @@ def test_get_transition_events_create_instance(topology: Topology, instance: Ins
def test_get_transition_events_delete_instance(topology: Topology, instance: Instance):
# arrange
instance_id = InstanceId()
- current_instances: dict[InstanceId, Instance] = {
- instance_id: instance
- }
+ current_instances: dict[InstanceId, Instance] = {instance_id: instance}
target_instances: dict[InstanceId, Instance] = {}
# act
diff --git a/src/exo/master/tests/test_placement_utils.py b/src/exo/master/tests/test_placement_utils.py
index 2e505779..ed1dadc2 100644
--- a/src/exo/master/tests/test_placement_utils.py
+++ b/src/exo/master/tests/test_placement_utils.py
@@ -21,23 +21,27 @@ def topology() -> Topology:
return topology
-def test_filter_cycles_by_memory(topology: Topology, create_node: Callable[[int, NodeId | None], Node], create_connection: Callable[[NodeId, NodeId], Connection]):
+def test_filter_cycles_by_memory(
+ topology: Topology,
+ create_node: Callable[[int, NodeId | None], Node],
+ create_connection: Callable[[NodeId, NodeId], Connection],
+):
# arrange
node1_id = NodeId()
node2_id = NodeId()
- node1 = create_node(1000*1024, node1_id)
- node2 = create_node(1000*1024, node2_id)
-
+ node1 = create_node(1000 * 1024, node1_id)
+ node2 = create_node(1000 * 1024, node2_id)
+
topology.add_node(node1)
topology.add_node(node2)
-
+
connection1 = create_connection(node1_id, node2_id)
connection2 = create_connection(node2_id, node1_id)
-
+
topology.add_connection(connection1)
topology.add_connection(connection2)
-
+
cycles = topology.get_cycles()
assert len(cycles) == 1
assert len(cycles[0]) == 2
@@ -51,69 +55,86 @@ def test_filter_cycles_by_memory(topology: Topology, create_node: Callable[[int,
assert set(n.node_id for n in filtered_cycles[0]) == {node1_id, node2_id}
-def test_filter_cycles_by_insufficient_memory(topology: Topology, create_node: Callable[[int, NodeId | None], Node], create_connection: Callable[[NodeId, NodeId], Connection]):
+def test_filter_cycles_by_insufficient_memory(
+ topology: Topology,
+ create_node: Callable[[int, NodeId | None], Node],
+ create_connection: Callable[[NodeId, NodeId], Connection],
+):
# arrange
node1_id = NodeId()
node2_id = NodeId()
- node1 = create_node(1000*1024, node1_id)
- node2 = create_node(1000*1024, node2_id)
+ node1 = create_node(1000 * 1024, node1_id)
+ node2 = create_node(1000 * 1024, node2_id)
topology.add_node(node1)
topology.add_node(node2)
connection1 = create_connection(node1_id, node2_id)
connection2 = create_connection(node2_id, node1_id)
-
+
topology.add_connection(connection1)
topology.add_connection(connection2)
-
+
# act
- filtered_cycles = filter_cycles_by_memory(topology.get_cycles(), 2001*1024)
+ filtered_cycles = filter_cycles_by_memory(topology.get_cycles(), 2001 * 1024)
# assert
assert len(filtered_cycles) == 0
-def test_filter_multiple_cycles_by_memory(topology: Topology, create_node: Callable[[int, NodeId | None], Node], create_connection: Callable[[NodeId, NodeId], Connection]):
+def test_filter_multiple_cycles_by_memory(
+ topology: Topology,
+ create_node: Callable[[int, NodeId | None], Node],
+ create_connection: Callable[[NodeId, NodeId], Connection],
+):
# arrange
node_a_id = NodeId()
node_b_id = NodeId()
node_c_id = NodeId()
-
- node_a = create_node(500*1024, node_a_id)
- node_b = create_node(500*1024, node_b_id)
- node_c = create_node(1000*1024, node_c_id)
-
+
+ node_a = create_node(500 * 1024, node_a_id)
+ node_b = create_node(500 * 1024, node_b_id)
+ node_c = create_node(1000 * 1024, node_c_id)
+
topology.add_node(node_a)
topology.add_node(node_b)
topology.add_node(node_c)
-
+
topology.add_connection(create_connection(node_a_id, node_b_id))
topology.add_connection(create_connection(node_b_id, node_a_id))
-
+
topology.add_connection(create_connection(node_a_id, node_c_id))
topology.add_connection(create_connection(node_c_id, node_b_id))
-
+
cycles = topology.get_cycles()
-
+
# act
- filtered_cycles = filter_cycles_by_memory(cycles, 1500*1024)
-
+ filtered_cycles = filter_cycles_by_memory(cycles, 1500 * 1024)
+
# assert
assert len(filtered_cycles) == 1
assert len(filtered_cycles[0]) == 3
- assert set(n.node_id for n in filtered_cycles[0]) == {node_a_id, node_b_id, node_c_id}
+ assert set(n.node_id for n in filtered_cycles[0]) == {
+ node_a_id,
+ node_b_id,
+ node_c_id,
+ }
-def test_get_smallest_cycles(topology: Topology, create_node: Callable[[int, NodeId | None], Node], create_connection: Callable[[NodeId, NodeId], Connection]):
+
+def test_get_smallest_cycles(
+ topology: Topology,
+ create_node: Callable[[int, NodeId | None], Node],
+ create_connection: Callable[[NodeId, NodeId], Connection],
+):
# arrange
node_a_id = NodeId()
node_b_id = NodeId()
node_c_id = NodeId()
-
- node_a = create_node(500*1024, node_a_id)
- node_b = create_node(500*1024, node_b_id)
- node_c = create_node(1000*1024, node_c_id)
+
+ node_a = create_node(500 * 1024, node_a_id)
+ node_b = create_node(500 * 1024, node_b_id)
+ node_c = create_node(1000 * 1024, node_c_id)
topology.add_node(node_a)
topology.add_node(node_b)
@@ -132,20 +153,31 @@ def test_get_smallest_cycles(topology: Topology, create_node: Callable[[int, Nod
assert len(smallest_cycles[0]) == 2
assert set(n.node_id for n in smallest_cycles[0]) == {node_a_id, node_b_id}
-@pytest.mark.parametrize("available_memory,total_layers,expected_layers", [
- ((500, 500, 1000), 12, (3, 3, 6)),
- ((500, 500, 500), 12, (4, 4, 4)),
- ((312, 518, 1024), 12, (2, 3, 7))
-])
-def test_get_shard_assignments(topology: Topology, create_node: Callable[[int, NodeId | None], Node], create_connection: Callable[[NodeId, NodeId], Connection], available_memory: tuple[int, int, int], total_layers: int, expected_layers: tuple[int, int, int]):
+
+@pytest.mark.parametrize(
+ "available_memory,total_layers,expected_layers",
+ [
+ ((500, 500, 1000), 12, (3, 3, 6)),
+ ((500, 500, 500), 12, (4, 4, 4)),
+ ((312, 518, 1024), 12, (2, 3, 7)),
+ ],
+)
+def test_get_shard_assignments(
+ topology: Topology,
+ create_node: Callable[[int, NodeId | None], Node],
+ create_connection: Callable[[NodeId, NodeId], Connection],
+ available_memory: tuple[int, int, int],
+ total_layers: int,
+ expected_layers: tuple[int, int, int],
+):
# arrange
node_a_id = NodeId()
node_b_id = NodeId()
node_c_id = NodeId()
-
- node_a = create_node(available_memory[0]*1024, node_a_id)
- node_b = create_node(available_memory[1]*1024, node_b_id)
- node_c = create_node(available_memory[2]*1024, node_c_id)
+
+ node_a = create_node(available_memory[0] * 1024, node_a_id)
+ node_b = create_node(available_memory[1] * 1024, node_b_id)
+ node_c = create_node(available_memory[2] * 1024, node_c_id)
topology.add_node(node_a)
topology.add_node(node_b)
@@ -155,16 +187,16 @@ def test_get_shard_assignments(topology: Topology, create_node: Callable[[int, N
topology.add_connection(create_connection(node_b_id, node_c_id))
topology.add_connection(create_connection(node_c_id, node_a_id))
topology.add_connection(create_connection(node_b_id, node_a_id))
-
+
model_meta = ModelMetadata(
model_id="test-model",
pretty_name="Test Model",
n_layers=total_layers,
- storage_size_kilobytes=1000
+ storage_size_kilobytes=1000,
)
cycles = topology.get_cycles()
selected_cycle = cycles[0]
-
+
# act
shard_assignments = get_shard_assignments(model_meta, selected_cycle)
@@ -172,25 +204,41 @@ def test_get_shard_assignments(topology: Topology, create_node: Callable[[int, N
runner_id_a = shard_assignments.node_to_runner[node_a_id]
runner_id_b = shard_assignments.node_to_runner[node_b_id]
runner_id_c = shard_assignments.node_to_runner[node_c_id]
- assert shard_assignments.runner_to_shard[runner_id_c].end_layer - shard_assignments.runner_to_shard[runner_id_c].start_layer == expected_layers[2]
- assert shard_assignments.runner_to_shard[runner_id_a].end_layer - shard_assignments.runner_to_shard[runner_id_a].start_layer == expected_layers[0]
- assert shard_assignments.runner_to_shard[runner_id_b].end_layer - shard_assignments.runner_to_shard[runner_id_b].start_layer == expected_layers[1]
+ assert (
+ shard_assignments.runner_to_shard[runner_id_c].end_layer
+ - shard_assignments.runner_to_shard[runner_id_c].start_layer
+ == expected_layers[2]
+ )
+ assert (
+ shard_assignments.runner_to_shard[runner_id_a].end_layer
+ - shard_assignments.runner_to_shard[runner_id_a].start_layer
+ == expected_layers[0]
+ )
+ assert (
+ shard_assignments.runner_to_shard[runner_id_b].end_layer
+ - shard_assignments.runner_to_shard[runner_id_b].start_layer
+ == expected_layers[1]
+ )
-def test_get_hosts_from_subgraph(topology: Topology, create_node: Callable[[int, NodeId | None], Node], create_connection: Callable[[NodeId, NodeId, int | None], Connection]):
+def test_get_hosts_from_subgraph(
+ topology: Topology,
+ create_node: Callable[[int, NodeId | None], Node],
+ create_connection: Callable[[NodeId, NodeId, int | None], Connection],
+):
# arrange
node_a_id = NodeId()
node_b_id = NodeId()
node_c_id = NodeId()
-
+
node_a = create_node(500, node_a_id)
node_b = create_node(500, node_b_id)
node_c = create_node(1000, node_c_id)
-
+
topology.add_node(node_a)
topology.add_node(node_b)
topology.add_node(node_c)
-
+
topology.add_connection(create_connection(node_a_id, node_b_id, 5001))
topology.add_connection(create_connection(node_b_id, node_c_id, 5002))
topology.add_connection(create_connection(node_c_id, node_a_id, 5003))
diff --git a/src/exo/master/tests/test_topology.py b/src/exo/master/tests/test_topology.py
index 32624723..18cb84a2 100644
--- a/src/exo/master/tests/test_topology.py
+++ b/src/exo/master/tests/test_topology.py
@@ -22,14 +22,26 @@ def connection() -> Connection:
send_back_node_id=NodeId(),
local_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/1234"),
send_back_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/1235"),
- connection_profile=ConnectionProfile(throughput=1000, latency=1000, jitter=1000))
+ connection_profile=ConnectionProfile(
+ throughput=1000, latency=1000, jitter=1000
+ ),
+ )
+
@pytest.fixture
def node_profile() -> NodePerformanceProfile:
- memory_profile = MemoryPerformanceProfile(ram_total=1000, ram_available=1000, swap_total=1000, swap_available=1000)
+ memory_profile = MemoryPerformanceProfile(
+ ram_total=1000, ram_available=1000, swap_total=1000, swap_available=1000
+ )
system_profile = SystemPerformanceProfile(flops_fp16=1000)
- return NodePerformanceProfile(model_id="test", chip_id="test", friendly_name="test", memory=memory_profile, network_interfaces=[],
- system=system_profile)
+ return NodePerformanceProfile(
+ model_id="test",
+ chip_id="test",
+ friendly_name="test",
+ memory=memory_profile,
+ network_interfaces=[],
+ system=system_profile,
+ )
@pytest.fixture
@@ -49,10 +61,14 @@ def test_add_node(topology: Topology, node_profile: NodePerformanceProfile):
assert data == node_profile
-def test_add_connection(topology: Topology, node_profile: NodePerformanceProfile, connection: Connection):
+def test_add_connection(
+ topology: Topology, node_profile: NodePerformanceProfile, connection: Connection
+):
# arrange
topology.add_node(Node(node_id=connection.local_node_id, node_profile=node_profile))
- topology.add_node(Node(node_id=connection.send_back_node_id, node_profile=node_profile))
+ topology.add_node(
+ Node(node_id=connection.send_back_node_id, node_profile=node_profile)
+ )
topology.add_connection(connection)
# act
@@ -62,38 +78,57 @@ def test_add_connection(topology: Topology, node_profile: NodePerformanceProfile
assert data == connection.connection_profile
-def test_update_node_profile(topology: Topology, node_profile: NodePerformanceProfile, connection: Connection):
+def test_update_node_profile(
+ topology: Topology, node_profile: NodePerformanceProfile, connection: Connection
+):
# arrange
topology.add_node(Node(node_id=connection.local_node_id, node_profile=node_profile))
- topology.add_node(Node(node_id=connection.send_back_node_id, node_profile=node_profile))
+ topology.add_node(
+ Node(node_id=connection.send_back_node_id, node_profile=node_profile)
+ )
topology.add_connection(connection)
- new_node_profile = NodePerformanceProfile(model_id="test", chip_id="test",
- friendly_name="test",
- memory=MemoryPerformanceProfile(ram_total=1000, ram_available=1000,
- swap_total=1000, swap_available=1000),
- network_interfaces=[],
- system=SystemPerformanceProfile(flops_fp16=1000))
+ new_node_profile = NodePerformanceProfile(
+ model_id="test",
+ chip_id="test",
+ friendly_name="test",
+ memory=MemoryPerformanceProfile(
+ ram_total=1000, ram_available=1000, swap_total=1000, swap_available=1000
+ ),
+ network_interfaces=[],
+ system=SystemPerformanceProfile(flops_fp16=1000),
+ )
# act
- topology.update_node_profile(connection.local_node_id, node_profile=new_node_profile)
+ topology.update_node_profile(
+ connection.local_node_id, node_profile=new_node_profile
+ )
# assert
data = topology.get_node_profile(connection.local_node_id)
assert data == new_node_profile
-def test_update_connection_profile(topology: Topology, node_profile: NodePerformanceProfile, connection: Connection):
+def test_update_connection_profile(
+ topology: Topology, node_profile: NodePerformanceProfile, connection: Connection
+):
# arrange
topology.add_node(Node(node_id=connection.local_node_id, node_profile=node_profile))
- topology.add_node(Node(node_id=connection.send_back_node_id, node_profile=node_profile))
+ topology.add_node(
+ Node(node_id=connection.send_back_node_id, node_profile=node_profile)
+ )
topology.add_connection(connection)
- new_connection_profile = ConnectionProfile(throughput=2000, latency=2000, jitter=2000)
- connection = Connection(local_node_id=connection.local_node_id, send_back_node_id=connection.send_back_node_id,
- local_multiaddr=connection.local_multiaddr,
- send_back_multiaddr=connection.send_back_multiaddr,
- connection_profile=new_connection_profile)
+ new_connection_profile = ConnectionProfile(
+ throughput=2000, latency=2000, jitter=2000
+ )
+ connection = Connection(
+ local_node_id=connection.local_node_id,
+ send_back_node_id=connection.send_back_node_id,
+ local_multiaddr=connection.local_multiaddr,
+ send_back_multiaddr=connection.send_back_multiaddr,
+ connection_profile=new_connection_profile,
+ )
# act
topology.update_connection_profile(connection)
@@ -103,11 +138,14 @@ def test_update_connection_profile(topology: Topology, node_profile: NodePerform
assert data == new_connection_profile
-def test_remove_connection_still_connected(topology: Topology, node_profile: NodePerformanceProfile,
- connection: Connection):
+def test_remove_connection_still_connected(
+ topology: Topology, node_profile: NodePerformanceProfile, connection: Connection
+):
# arrange
topology.add_node(Node(node_id=connection.local_node_id, node_profile=node_profile))
- topology.add_node(Node(node_id=connection.send_back_node_id, node_profile=node_profile))
+ topology.add_node(
+ Node(node_id=connection.send_back_node_id, node_profile=node_profile)
+ )
topology.add_connection(connection)
# act
@@ -117,7 +155,9 @@ def test_remove_connection_still_connected(topology: Topology, node_profile: Nod
assert topology.get_connection_profile(connection) is None
-def test_remove_connection_bridge(topology: Topology, node_profile: NodePerformanceProfile, connection: Connection):
+def test_remove_connection_bridge(
+ topology: Topology, node_profile: NodePerformanceProfile, connection: Connection
+):
"""Create a bridge scenario: master -> node_a -> node_b
and remove the bridge connection (master -> node_a)"""
# arrange
@@ -128,15 +168,17 @@ def test_remove_connection_bridge(topology: Topology, node_profile: NodePerforma
topology.add_node(Node(node_id=master_id, node_profile=node_profile))
topology.add_node(Node(node_id=node_a_id, node_profile=node_profile))
topology.add_node(Node(node_id=node_b_id, node_profile=node_profile))
-
+
topology.set_master_node_id(master_id)
-
+
connection_master_to_a = Connection(
local_node_id=master_id,
send_back_node_id=node_a_id,
local_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/1234"),
send_back_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/1235"),
- connection_profile=ConnectionProfile(throughput=1000, latency=1000, jitter=1000)
+ connection_profile=ConnectionProfile(
+ throughput=1000, latency=1000, jitter=1000
+ ),
)
connection_a_to_b = Connection(
@@ -144,7 +186,9 @@ def test_remove_connection_bridge(topology: Topology, node_profile: NodePerforma
send_back_node_id=node_b_id,
local_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/1236"),
send_back_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/1237"),
- connection_profile=ConnectionProfile(throughput=1000, latency=1000, jitter=1000)
+ connection_profile=ConnectionProfile(
+ throughput=1000, latency=1000, jitter=1000
+ ),
)
topology.add_connection(connection_master_to_a)
@@ -162,10 +206,14 @@ def test_remove_connection_bridge(topology: Topology, node_profile: NodePerforma
assert topology.get_node_profile(node_b_id) is None
-def test_remove_node_still_connected(topology: Topology, node_profile: NodePerformanceProfile, connection: Connection):
+def test_remove_node_still_connected(
+ topology: Topology, node_profile: NodePerformanceProfile, connection: Connection
+):
# arrange
topology.add_node(Node(node_id=connection.local_node_id, node_profile=node_profile))
- topology.add_node(Node(node_id=connection.send_back_node_id, node_profile=node_profile))
+ topology.add_node(
+ Node(node_id=connection.send_back_node_id, node_profile=node_profile)
+ )
topology.add_connection(connection)
# act
@@ -175,10 +223,14 @@ def test_remove_node_still_connected(topology: Topology, node_profile: NodePerfo
assert topology.get_node_profile(connection.local_node_id) is None
-def test_list_nodes(topology: Topology, node_profile: NodePerformanceProfile, connection: Connection):
+def test_list_nodes(
+ topology: Topology, node_profile: NodePerformanceProfile, connection: Connection
+):
# arrange
topology.add_node(Node(node_id=connection.local_node_id, node_profile=node_profile))
- topology.add_node(Node(node_id=connection.send_back_node_id, node_profile=node_profile))
+ topology.add_node(
+ Node(node_id=connection.send_back_node_id, node_profile=node_profile)
+ )
topology.add_connection(connection)
# act
@@ -187,4 +239,7 @@ def test_list_nodes(topology: Topology, node_profile: NodePerformanceProfile, co
# assert
assert len(nodes) == 2
assert all(isinstance(node, Node) for node in nodes)
- assert {node.node_id for node in nodes} == {connection.local_node_id, connection.send_back_node_id}
+ assert {node.node_id for node in nodes} == {
+ connection.local_node_id,
+ connection.send_back_node_id,
+ }
diff --git a/src/exo/master/utils/placement_utils.py b/src/exo/master/utils/placement_utils.py
index 86cf14d2..b89736b1 100644
--- a/src/exo/master/utils/placement_utils.py
+++ b/src/exo/master/utils/placement_utils.py
@@ -16,10 +16,14 @@ class NodeWithProfile(BaseModel):
node_id: NodeId
node_profile: NodePerformanceProfile
+
def narrow_all_nodes(nodes: list[Node]) -> TypeGuard[list[NodeWithProfile]]:
return all(node.node_profile is not None for node in nodes)
-def filter_cycles_by_memory(cycles: list[list[Node]], required_memory: int) -> list[list[Node]]:
+
+def filter_cycles_by_memory(
+ cycles: list[list[Node]], required_memory: int
+) -> list[list[Node]]:
filtered_cycles: list[list[Node]] = []
for cycle in cycles:
if not narrow_all_nodes(cycle):
@@ -35,6 +39,7 @@ def get_smallest_cycles(cycles: list[list[Node]]) -> list[list[Node]]:
min_nodes = min(len(cycle) for cycle in cycles)
return [cycle for cycle in cycles if len(cycle) == min_nodes]
+
def get_shard_assignments(
model_meta: ModelMetadata,
selected_cycle: list[Node],
@@ -42,7 +47,9 @@ def get_shard_assignments(
if not narrow_all_nodes(selected_cycle):
raise ValueError("All nodes must have profiles to create shard assignments")
- cycle_memory = sum(node.node_profile.memory.ram_available for node in selected_cycle)
+ cycle_memory = sum(
+ node.node_profile.memory.ram_available for node in selected_cycle
+ )
total_layers = model_meta.n_layers
runner_to_shard: dict[RunnerId, PipelineShardMetadata] = {}
node_to_runner: dict[NodeId, RunnerId] = {}
@@ -52,7 +59,9 @@ def get_shard_assignments(
if i == len(selected_cycle) - 1:
node_layers = total_layers - layers_assigned
else:
- node_layers = round(total_layers * (node.node_profile.memory.ram_available / cycle_memory))
+ node_layers = round(
+ total_layers * (node.node_profile.memory.ram_available / cycle_memory)
+ )
node_layers = max(1, node_layers)
runner_id = RunnerId()
@@ -62,7 +71,7 @@ def get_shard_assignments(
world_size=len(selected_cycle),
start_layer=layers_assigned,
end_layer=layers_assigned + node_layers,
- n_layers=total_layers
+ n_layers=total_layers,
)
runner_to_shard[runner_id] = shard
@@ -72,7 +81,7 @@ def get_shard_assignments(
shard_assignments = ShardAssignments(
model_id=model_meta.model_id,
runner_to_shard=runner_to_shard,
- node_to_runner=node_to_runner
+ node_to_runner=node_to_runner,
)
return shard_assignments
@@ -82,27 +91,29 @@ def get_hosts_from_subgraph(cycle_digraph: Topology) -> list[Host]:
cycles = cycle_digraph.get_cycles()
if not cycles:
return []
-
+
get_thunderbolt = False
if cycle_digraph.is_thunderbolt_cycle(cycles[0]):
get_thunderbolt = True
-
+
cycle = cycles[0]
hosts: list[Host] = []
for i in range(len(cycle)):
current_node = cycle[i]
next_node = cycle[(i + 1) % len(cycle)]
-
+
for connection in cycle_digraph.list_connections():
- if (connection.local_node_id == current_node.node_id and
- connection.send_back_node_id == next_node.node_id):
+ if (
+ connection.local_node_id == current_node.node_id
+ and connection.send_back_node_id == next_node.node_id
+ ):
if get_thunderbolt and not connection.is_thunderbolt():
continue
host = Host(
ip=connection.send_back_multiaddr.ip_address,
- port=connection.send_back_multiaddr.port
+ port=connection.send_back_multiaddr.port,
)
hosts.append(host)
break
-
- return hosts
\ No newline at end of file
+
+ return hosts
diff --git a/src/exo/shared/__init__.py b/src/exo/shared/__init__.py
index 0519ecba..e69de29b 100644
--- a/src/exo/shared/__init__.py
+++ b/src/exo/shared/__init__.py
@@ -1 +0,0 @@
-
\ No newline at end of file
diff --git a/src/exo/shared/apply/__init__.py b/src/exo/shared/apply/__init__.py
index 534e5356..dc22de1e 100644
--- a/src/exo/shared/apply/__init__.py
+++ b/src/exo/shared/apply/__init__.py
@@ -1,3 +1,3 @@
from .apply import apply
-__all__ = ["apply"]
\ No newline at end of file
+__all__ = ["apply"]
diff --git a/src/exo/shared/apply/apply.py b/src/exo/shared/apply/apply.py
index 134ce3c8..75c102f4 100644
--- a/src/exo/shared/apply/apply.py
+++ b/src/exo/shared/apply/apply.py
@@ -49,25 +49,31 @@ def event_apply(event: Event, state: State) -> State:
raise RuntimeError(f"no handler registered for event type {type(event).__name__}")
+
def apply(state: State, event: EventFromEventLog[Event]) -> State:
new_state: State = event_apply(event.event, state)
return new_state.model_copy(update={"last_event_applied_idx": event.idx_in_log})
+
@event_apply.register(TaskCreated)
def apply_task_created(event: TaskCreated, state: State) -> State:
new_tasks: Mapping[TaskId, Task] = {**state.tasks, event.task_id: event.task}
return state.model_copy(update={"tasks": new_tasks})
+
@event_apply.register(TaskDeleted)
def apply_task_deleted(event: TaskDeleted, state: State) -> State:
- new_tasks: Mapping[TaskId, Task] = {tid: task for tid, task in state.tasks.items() if tid != event.task_id}
+ new_tasks: Mapping[TaskId, Task] = {
+ tid: task for tid, task in state.tasks.items() if tid != event.task_id
+ }
return state.model_copy(update={"tasks": new_tasks})
+
@event_apply.register(TaskStateUpdated)
def apply_task_state_updated(event: TaskStateUpdated, state: State) -> State:
if event.task_id not in state.tasks:
return state
-
+
update: dict[str, TaskStatus | None] = {
"task_status": event.task_status,
}
@@ -79,46 +85,71 @@ def apply_task_state_updated(event: TaskStateUpdated, state: State) -> State:
new_tasks: Mapping[TaskId, Task] = {**state.tasks, event.task_id: updated_task}
return state.model_copy(update={"tasks": new_tasks})
+
@event_apply.register(TaskFailed)
def apply_task_failed(event: TaskFailed, state: State) -> State:
if event.task_id not in state.tasks:
return state
-
- updated_task = state.tasks[event.task_id].model_copy(update={"error_type": event.error_type, "error_message": event.error_message})
+
+ updated_task = state.tasks[event.task_id].model_copy(
+ update={"error_type": event.error_type, "error_message": event.error_message}
+ )
new_tasks: Mapping[TaskId, Task] = {**state.tasks, event.task_id: updated_task}
return state.model_copy(update={"tasks": new_tasks})
+
@event_apply.register(InstanceCreated)
def apply_instance_created(event: InstanceCreated, state: State) -> State:
instance = event.instance
- new_instances: Mapping[InstanceId, Instance] = {**state.instances, instance.instance_id: instance}
+ new_instances: Mapping[InstanceId, Instance] = {
+ **state.instances,
+ instance.instance_id: instance,
+ }
return state.model_copy(update={"instances": new_instances})
+
@event_apply.register(InstanceActivated)
def apply_instance_activated(event: InstanceActivated, state: State) -> State:
if event.instance_id not in state.instances:
return state
-
- updated_instance = state.instances[event.instance_id].model_copy(update={"type": InstanceStatus.ACTIVE})
- new_instances: Mapping[InstanceId, Instance] = {**state.instances, event.instance_id: updated_instance}
+
+ updated_instance = state.instances[event.instance_id].model_copy(
+ update={"type": InstanceStatus.ACTIVE}
+ )
+ new_instances: Mapping[InstanceId, Instance] = {
+ **state.instances,
+ event.instance_id: updated_instance,
+ }
return state.model_copy(update={"instances": new_instances})
+
@event_apply.register(InstanceDeactivated)
def apply_instance_deactivated(event: InstanceDeactivated, state: State) -> State:
if event.instance_id not in state.instances:
return state
-
- updated_instance = state.instances[event.instance_id].model_copy(update={"type": InstanceStatus.INACTIVE})
- new_instances: Mapping[InstanceId, Instance] = {**state.instances, event.instance_id: updated_instance}
+
+ updated_instance = state.instances[event.instance_id].model_copy(
+ update={"type": InstanceStatus.INACTIVE}
+ )
+ new_instances: Mapping[InstanceId, Instance] = {
+ **state.instances,
+ event.instance_id: updated_instance,
+ }
return state.model_copy(update={"instances": new_instances})
+
@event_apply.register(InstanceDeleted)
def apply_instance_deleted(event: InstanceDeleted, state: State) -> State:
- new_instances: Mapping[InstanceId, Instance] = {iid: inst for iid, inst in state.instances.items() if iid != event.instance_id}
+ new_instances: Mapping[InstanceId, Instance] = {
+ iid: inst for iid, inst in state.instances.items() if iid != event.instance_id
+ }
return state.model_copy(update={"instances": new_instances})
+
@event_apply.register(InstanceReplacedAtomically)
-def apply_instance_replaced_atomically(event: InstanceReplacedAtomically, state: State) -> State:
+def apply_instance_replaced_atomically(
+ event: InstanceReplacedAtomically, state: State
+) -> State:
new_instances = dict(state.instances)
if event.instance_to_replace in new_instances:
del new_instances[event.instance_to_replace]
@@ -126,19 +157,32 @@ def apply_instance_replaced_atomically(event: InstanceReplacedAtomically, state:
new_instances[event.new_instance_id] = state.instances[event.new_instance_id]
return state.model_copy(update={"instances": new_instances})
+
@event_apply.register(RunnerStatusUpdated)
def apply_runner_status_updated(event: RunnerStatusUpdated, state: State) -> State:
- new_runners: Mapping[RunnerId, RunnerStatus] = {**state.runners, event.runner_id: event.runner_status}
+ new_runners: Mapping[RunnerId, RunnerStatus] = {
+ **state.runners,
+ event.runner_id: event.runner_status,
+ }
return state.model_copy(update={"runners": new_runners})
+
@event_apply.register(RunnerDeleted)
def apply_runner_deleted(event: RunnerDeleted, state: State) -> State:
- new_runners: Mapping[RunnerId, RunnerStatus] = {rid: rs for rid, rs in state.runners.items() if rid != event.runner_id}
+ new_runners: Mapping[RunnerId, RunnerStatus] = {
+ rid: rs for rid, rs in state.runners.items() if rid != event.runner_id
+ }
return state.model_copy(update={"runners": new_runners})
+
@event_apply.register(NodePerformanceMeasured)
-def apply_node_performance_measured(event: NodePerformanceMeasured, state: State) -> State:
- new_profiles: Mapping[NodeId, NodePerformanceProfile] = {**state.node_profiles, event.node_id: event.node_profile}
+def apply_node_performance_measured(
+ event: NodePerformanceMeasured, state: State
+) -> State:
+ new_profiles: Mapping[NodeId, NodePerformanceProfile] = {
+ **state.node_profiles,
+ event.node_id: event.node_profile,
+ }
state = state.model_copy(update={"node_profiles": new_profiles})
topology = copy.copy(state.topology)
if not topology.contains_node(event.node_id):
@@ -147,11 +191,16 @@ def apply_node_performance_measured(event: NodePerformanceMeasured, state: State
topology.update_node_profile(event.node_id, event.node_profile)
return state.model_copy(update={"topology": topology})
+
@event_apply.register(WorkerStatusUpdated)
def apply_worker_status_updated(event: WorkerStatusUpdated, state: State) -> State:
- new_node_status: Mapping[NodeId, NodeStatus] = {**state.node_status, event.node_id: event.node_state}
+ new_node_status: Mapping[NodeId, NodeStatus] = {
+ **state.node_status,
+ event.node_id: event.node_state,
+ }
return state.model_copy(update={"node_status": new_node_status})
+
@event_apply.register(TopologyNodeCreated)
def apply_topology_node_created(event: TopologyNodeCreated, state: State) -> State:
topology = copy.copy(state.topology)
@@ -160,18 +209,23 @@ def apply_topology_node_created(event: TopologyNodeCreated, state: State) -> Sta
topology.set_master_node_id(event.node_id)
return state.model_copy(update={"topology": topology})
+
@event_apply.register(TopologyEdgeCreated)
def apply_topology_edge_created(event: TopologyEdgeCreated, state: State) -> State:
topology = copy.copy(state.topology)
topology.add_connection(event.edge)
return state.model_copy(update={"topology": topology})
+
@event_apply.register(TopologyEdgeReplacedAtomically)
-def apply_topology_edge_replaced_atomically(event: TopologyEdgeReplacedAtomically, state: State) -> State:
+def apply_topology_edge_replaced_atomically(
+ event: TopologyEdgeReplacedAtomically, state: State
+) -> State:
topology = copy.copy(state.topology)
topology.update_connection_profile(event.edge)
return state.model_copy(update={"topology": topology})
+
@event_apply.register(TopologyEdgeDeleted)
def apply_topology_edge_deleted(event: TopologyEdgeDeleted, state: State) -> State:
topology = copy.copy(state.topology)
@@ -182,9 +236,9 @@ def apply_topology_edge_deleted(event: TopologyEdgeDeleted, state: State) -> Sta
local_node_id=event.edge.send_back_node_id,
send_back_node_id=event.edge.local_node_id,
local_multiaddr=event.edge.send_back_multiaddr,
- send_back_multiaddr=event.edge.local_multiaddr
+ send_back_multiaddr=event.edge.local_multiaddr,
)
if not topology.contains_connection(opposite_edge):
return state.model_copy(update={"topology": topology})
topology.remove_connection(opposite_edge)
- return state.model_copy(update={"topology": topology})
\ No newline at end of file
+ return state.model_copy(update={"topology": topology})
diff --git a/src/exo/shared/constants.py b/src/exo/shared/constants.py
index acd0f569..fe1393c3 100644
--- a/src/exo/shared/constants.py
+++ b/src/exo/shared/constants.py
@@ -25,6 +25,7 @@ LB_TFLOPS = 2.3
LB_MEMBW_GBPS = 68
LB_DISK_GBPS = 1.5
+
# little helper function to get the name of the module that raised the error
def get_caller_module_name() -> str:
frm = inspect.stack()[1]
diff --git a/src/exo/shared/db/__init__.py b/src/exo/shared/db/__init__.py
index f7eb8bbc..955a46e2 100644
--- a/src/exo/shared/db/__init__.py
+++ b/src/exo/shared/db/__init__.py
@@ -2,4 +2,4 @@
from .sqlite import AsyncSQLiteEventStorage, EventStorageProtocol
-__all__ = ["AsyncSQLiteEventStorage", "EventStorageProtocol"]
\ No newline at end of file
+__all__ = ["AsyncSQLiteEventStorage", "EventStorageProtocol"]
diff --git a/src/exo/shared/db/sqlite/__init__.py b/src/exo/shared/db/sqlite/__init__.py
index abf926ff..d6c08ef5 100644
--- a/src/exo/shared/db/sqlite/__init__.py
+++ b/src/exo/shared/db/sqlite/__init__.py
@@ -12,4 +12,4 @@ __all__ = [
"EventLogType",
"EventStorageProtocol",
"StoredEvent",
-]
\ No newline at end of file
+]
diff --git a/src/exo/shared/db/sqlite/config.py b/src/exo/shared/db/sqlite/config.py
index dda4753a..f6f6ac97 100644
--- a/src/exo/shared/db/sqlite/config.py
+++ b/src/exo/shared/db/sqlite/config.py
@@ -8,19 +8,20 @@ from exo.shared.constants import EXO_GLOBAL_EVENT_DB, EXO_WORKER_EVENT_DB
class EventLogType(str, Enum):
"""Types of event logs in the system"""
+
WORKER_EVENTS = "worker_events"
GLOBAL_EVENTS = "global_events"
class EventLogConfig(BaseModel):
"""Configuration for the event log system"""
-
+
# Batch processing settings
batch_size: int = 100
batch_timeout_ms: int = 100
debounce_ms: int = 10
max_age_ms: int = 100
-
+
def get_db_path(self, log_type: EventLogType) -> Path:
"""Get the full path for a specific event log type"""
if log_type == EventLogType.WORKER_EVENTS:
@@ -28,4 +29,4 @@ class EventLogConfig(BaseModel):
elif log_type == EventLogType.GLOBAL_EVENTS:
return EXO_GLOBAL_EVENT_DB
else:
- raise ValueError(f"Unknown log type: {log_type}")
\ No newline at end of file
+ raise ValueError(f"Unknown log type: {log_type}")
diff --git a/src/exo/shared/db/sqlite/connector.py b/src/exo/shared/db/sqlite/connector.py
index e5b9793d..7a6d0767 100644
--- a/src/exo/shared/db/sqlite/connector.py
+++ b/src/exo/shared/db/sqlite/connector.py
@@ -21,27 +21,27 @@ from .types import StoredEvent
class AsyncSQLiteEventStorage:
"""High-performance SQLite event storage with async batching.
-
+
Features:
- Non-blocking writes via adaptive async batching with debouncing
- Automatic sequence numbering using SQLite rowid
- Type-safe event serialization/deserialization
- Efficient indexing for common query patterns
-
+
Batching behavior:
- Low load: Minimal latency via short debounce windows
- High load: Efficient batching up to batch_size limit
- Max age constraint prevents indefinite delays
"""
-
+
def __init__(
- self,
- db_path: str | Path,
+ self,
+ db_path: str | Path,
batch_size: int,
batch_timeout_ms: int,
debounce_ms: int,
max_age_ms: int,
- logger: Logger | None = None
+ logger: Logger | None = None,
):
self._db_path = Path(db_path)
self._batch_size = batch_size
@@ -49,56 +49,52 @@ class AsyncSQLiteEventStorage:
self._debounce_s = debounce_ms / 1000.0
self._max_age_s = max_age_ms / 1000.0
self._logger = logger or getLogger(__name__)
-
+
self._write_queue: Queue[tuple[Event, NodeId]] = Queue()
self._batch_writer_task: Task[None] | None = None
self._engine = None
self._closed = False
-
+
async def start(self) -> None:
"""Initialize the storage and start the batch writer."""
if self._batch_writer_task is not None:
raise RuntimeError("Storage already started")
-
+
# Create database and tables
await self._initialize_database()
-
+
# Start batch writer
self._batch_writer_task = asyncio.create_task(self._batch_writer())
self._logger.info(f"Started SQLite event storage: {self._db_path}")
-
- async def append_events(
- self,
- events: Sequence[Event],
- origin: NodeId
- ) -> None:
- """Append events to the log (fire-and-forget). The writes are batched and committed
+
+ async def append_events(self, events: Sequence[Event], origin: NodeId) -> None:
+ """Append events to the log (fire-and-forget). The writes are batched and committed
in the background so readers don't have a guarantee of seeing events immediately."""
if self._closed:
raise RuntimeError("Storage is closed")
-
+
for event in events:
await self._write_queue.put((event, origin))
-
+
async def get_events_since(
- self,
- last_idx: int,
- ignore_no_op_events: bool = False
+ self, last_idx: int, ignore_no_op_events: bool = False
) -> Sequence[EventFromEventLog[Event]]:
"""Retrieve events after a specific index."""
if self._closed:
raise RuntimeError("Storage is closed")
-
+
assert self._engine is not None
-
+
async with AsyncSession(self._engine) as session:
# Use raw SQL to get rowid along with the stored event data
result = await session.execute(
- text("SELECT rowid, origin, event_data FROM events WHERE rowid > :last_idx ORDER BY rowid"),
- {"last_idx": last_idx}
+ text(
+ "SELECT rowid, origin, event_data FROM events WHERE rowid > :last_idx ORDER BY rowid"
+ ),
+ {"last_idx": last_idx},
)
rows = result.fetchall()
-
+
events: list[EventFromEventLog[Event]] = []
for row in rows:
rowid: int = cast(int, row[0])
@@ -106,30 +102,36 @@ class AsyncSQLiteEventStorage:
# Parse JSON string to dict
raw_event_data = row[2] # type: ignore[reportAny] - SQLAlchemy result is Any
if isinstance(raw_event_data, str):
- event_data: dict[str, Any] = cast(dict[str, Any], json.loads(raw_event_data))
+ event_data: dict[str, Any] = cast(
+ dict[str, Any], json.loads(raw_event_data)
+ )
else:
event_data = cast(dict[str, Any], raw_event_data)
event = EventParser.validate_python(event_data)
if ignore_no_op_events and event.__no_apply__:
continue
- events.append(EventFromEventLog(
- event=event,
- origin=NodeId(origin),
- idx_in_log=rowid # rowid becomes idx_in_log
- ))
-
+ events.append(
+ EventFromEventLog(
+ event=event,
+ origin=NodeId(origin),
+ idx_in_log=rowid, # rowid becomes idx_in_log
+ )
+ )
+
return events
async def get_last_idx(self) -> int:
if self._closed:
raise RuntimeError("Storaged is closed")
-
+
assert self._engine is not None
async with AsyncSession(self._engine) as session:
result = await session.execute(
- text("SELECT rowid, origin, event_data FROM events ORDER BY rowid DESC LIMIT 1"),
- {}
+ text(
+ "SELECT rowid, origin, event_data FROM events ORDER BY rowid DESC LIMIT 1"
+ ),
+ {},
)
rows = result.fetchall()
@@ -139,34 +141,36 @@ class AsyncSQLiteEventStorage:
row = rows[0]
return cast(int, row[0])
else:
- raise AssertionError("There should have been at most 1 row returned from this SQL query.")
-
+ raise AssertionError(
+ "There should have been at most 1 row returned from this SQL query."
+ )
+
async def close(self) -> None:
"""Close the storage connection and cleanup resources."""
if self._closed:
return
-
+
self._closed = True
-
+
# Stop batch writer
if self._batch_writer_task is not None:
self._batch_writer_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await self._batch_writer_task
-
+
# Close database
if self._engine is not None:
await self._engine.dispose()
-
+
self._logger.info("Closed SQLite event storage")
-
+
async def delete_all_events(self) -> None:
"""Delete all events from the database."""
assert self._engine is not None
async with AsyncSession(self._engine) as session:
await session.execute(text("DELETE FROM events"))
await session.commit()
-
+
async def _initialize_database(self) -> None:
"""Initialize database connection and create tables."""
self._engine = create_async_engine(
@@ -178,22 +182,25 @@ class AsyncSQLiteEventStorage:
},
pool_pre_ping=True, # Test connections before using them
pool_size=5,
- max_overflow=10
+ max_overflow=10,
)
-
+
# Create tables with proper race condition handling
async with self._engine.begin() as conn:
# First check if the table exists using SQLite's master table
result = await conn.execute(
- text("SELECT name FROM sqlite_master WHERE type='table' AND name='events'")
+ text(
+ "SELECT name FROM sqlite_master WHERE type='table' AND name='events'"
+ )
)
table_exists = result.fetchone() is not None
-
+
if not table_exists:
try:
# Use CREATE TABLE IF NOT EXISTS as a more atomic operation
# This avoids race conditions between check and create
- await conn.execute(text("""
+ await conn.execute(
+ text("""
CREATE TABLE IF NOT EXISTS events (
rowid INTEGER PRIMARY KEY AUTOINCREMENT,
origin TEXT NOT NULL,
@@ -202,41 +209,69 @@ class AsyncSQLiteEventStorage:
event_data TEXT NOT NULL,
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
)
- """))
-
+ """)
+ )
+
# Create indexes if they don't exist
- await conn.execute(text("CREATE INDEX IF NOT EXISTS idx_events_origin ON events(origin)"))
- await conn.execute(text("CREATE INDEX IF NOT EXISTS idx_events_event_type ON events(event_type)"))
- await conn.execute(text("CREATE INDEX IF NOT EXISTS idx_events_event_id ON events(event_id)"))
- await conn.execute(text("CREATE INDEX IF NOT EXISTS idx_events_created_at ON events(created_at)"))
- await conn.execute(text("CREATE INDEX IF NOT EXISTS idx_events_origin_created ON events(origin, created_at)"))
-
+ await conn.execute(
+ text(
+ "CREATE INDEX IF NOT EXISTS idx_events_origin ON events(origin)"
+ )
+ )
+ await conn.execute(
+ text(
+ "CREATE INDEX IF NOT EXISTS idx_events_event_type ON events(event_type)"
+ )
+ )
+ await conn.execute(
+ text(
+ "CREATE INDEX IF NOT EXISTS idx_events_event_id ON events(event_id)"
+ )
+ )
+ await conn.execute(
+ text(
+ "CREATE INDEX IF NOT EXISTS idx_events_created_at ON events(created_at)"
+ )
+ )
+ await conn.execute(
+ text(
+ "CREATE INDEX IF NOT EXISTS idx_events_origin_created ON events(origin, created_at)"
+ )
+ )
+
self._logger.info("Events table and indexes created successfully")
except OperationalError as e:
# Even with IF NOT EXISTS, log any unexpected errors
self._logger.error(f"Error creating table: {e}")
# Re-check if table exists now
result = await conn.execute(
- text("SELECT name FROM sqlite_master WHERE type='table' AND name='events'")
+ text(
+ "SELECT name FROM sqlite_master WHERE type='table' AND name='events'"
+ )
)
if result.fetchone() is None:
raise RuntimeError(f"Failed to create events table: {e}") from e
else:
- self._logger.info("Events table exists (likely created by another process)")
+ self._logger.info(
+ "Events table exists (likely created by another process)"
+ )
else:
self._logger.debug("Events table already exists")
-
+
# Enable WAL mode and other optimizations with retry logic
- await self._execute_pragma_with_retry(conn, [
- "PRAGMA journal_mode=WAL",
- "PRAGMA synchronous=NORMAL",
- "PRAGMA cache_size=10000",
- "PRAGMA busy_timeout=30000" # 30 seconds busy timeout
- ])
-
+ await self._execute_pragma_with_retry(
+ conn,
+ [
+ "PRAGMA journal_mode=WAL",
+ "PRAGMA synchronous=NORMAL",
+ "PRAGMA cache_size=10000",
+ "PRAGMA busy_timeout=30000", # 30 seconds busy timeout
+ ],
+ )
+
async def _batch_writer(self) -> None:
"""Background task that drains the queue and commits batches.
-
+
Uses adaptive batching with debouncing:
- Blocks waiting for first item (no CPU waste when idle)
- Opens debounce window to collect more items
@@ -244,50 +279,50 @@ class AsyncSQLiteEventStorage:
- Resets debounce timer with each new item
"""
loop = asyncio.get_event_loop()
-
+
while not self._closed:
batch: list[tuple[Event, NodeId]] = []
-
+
try:
# Block waiting for first item
event, origin = await self._write_queue.get()
batch.append((event, origin))
first_ts = loop.time() # monotonic seconds
-
+
# Open debounce window
while True:
# How much longer can we wait?
age_left = self._max_age_s - (loop.time() - first_ts)
if age_left <= 0:
break # max age reached → flush
-
+
# Shrink the wait to honour both debounce and max-age
try:
event, origin = await asyncio.wait_for(
self._write_queue.get(),
- timeout=min(self._debounce_s, age_left)
+ timeout=min(self._debounce_s, age_left),
)
batch.append((event, origin))
-
+
if len(batch) >= self._batch_size:
break # size cap reached → flush
# else: loop again, resetting debounce timer
except asyncio.TimeoutError:
break # debounce window closed → flush
-
+
except asyncio.CancelledError:
# Drain any remaining items before exiting
if batch:
await self._commit_batch(batch)
raise
-
+
if batch:
await self._commit_batch(batch)
-
+
async def _commit_batch(self, batch: list[tuple[Event, NodeId]]) -> None:
"""Commit a batch of events to SQLite."""
assert self._engine is not None
-
+
try:
async with AsyncSession(self._engine) as session:
for event, origin in batch:
@@ -295,17 +330,21 @@ class AsyncSQLiteEventStorage:
origin=origin,
event_type=event.event_type,
event_id=str(event.event_id),
- event_data=event.model_dump(mode='json') # Serialize UUIDs and other objects to JSON-compatible strings
+ event_data=event.model_dump(
+ mode="json"
+ ), # Serialize UUIDs and other objects to JSON-compatible strings
)
session.add(stored_event)
-
+
await session.commit()
if len([ev for ev in batch if not isinstance(ev[0], Heartbeat)]) > 0:
self._logger.debug(f"Committed batch of {len(batch)} events")
-
+
except OperationalError as e:
if "database is locked" in str(e):
- self._logger.warning(f"Database locked during batch commit, will retry: {e}")
+ self._logger.warning(
+ f"Database locked during batch commit, will retry: {e}"
+ )
# Retry with exponential backoff
await self._commit_batch_with_retry(batch)
else:
@@ -314,58 +353,77 @@ class AsyncSQLiteEventStorage:
except Exception as e:
self._logger.error(f"Failed to commit batch: {e}")
raise
-
- async def _execute_pragma_with_retry(self, conn: AsyncConnection, pragmas: list[str], max_retries: int = 5) -> None:
+
+ async def _execute_pragma_with_retry(
+ self, conn: AsyncConnection, pragmas: list[str], max_retries: int = 5
+ ) -> None:
"""Execute PRAGMA statements with retry logic for database lock errors."""
for pragma in pragmas:
retry_count = 0
base_delay: float = 0.1 # 100ms
-
+
while retry_count < max_retries:
try:
await conn.execute(text(pragma))
break
except OperationalError as e:
if "database is locked" in str(e) and retry_count < max_retries - 1:
- delay = cast(float, base_delay * (2 ** retry_count) + random.uniform(0, 0.1))
- self._logger.warning(f"Database locked on '{pragma}', retry {retry_count + 1}/{max_retries} after {delay:.2f}s")
+ delay = cast(
+ float,
+ base_delay * (2**retry_count) + random.uniform(0, 0.1),
+ )
+ self._logger.warning(
+ f"Database locked on '{pragma}', retry {retry_count + 1}/{max_retries} after {delay:.2f}s"
+ )
await asyncio.sleep(delay)
retry_count += 1
else:
- self._logger.error(f"Failed to execute '{pragma}' after {retry_count + 1} attempts: {e}")
+ self._logger.error(
+ f"Failed to execute '{pragma}' after {retry_count + 1} attempts: {e}"
+ )
raise
-
- async def _commit_batch_with_retry(self, batch: list[tuple[Event, NodeId]], max_retries: int = 5) -> None:
+
+ async def _commit_batch_with_retry(
+ self, batch: list[tuple[Event, NodeId]], max_retries: int = 5
+ ) -> None:
"""Commit a batch with retry logic for database lock errors."""
retry_count = 0
base_delay: float = 0.1 # 100ms
-
+
while retry_count < max_retries:
try:
assert self._engine is not None
-
+
async with AsyncSession(self._engine) as session:
for event, origin in batch:
stored_event = StoredEvent(
origin=origin,
event_type=event.event_type,
event_id=str(event.event_id),
- event_data=event.model_dump(mode='json')
+ event_data=event.model_dump(mode="json"),
)
session.add(stored_event)
-
+
await session.commit()
-
+
if len([ev for ev in batch if not isinstance(ev[0], Heartbeat)]) > 0:
- self._logger.debug(f"Committed batch of {len(batch)} events after {retry_count} retries")
+ self._logger.debug(
+ f"Committed batch of {len(batch)} events after {retry_count} retries"
+ )
return
-
+
except OperationalError as e:
if "database is locked" in str(e) and retry_count < max_retries - 1:
- delay = cast(float, base_delay * (2 ** retry_count) + random.uniform(0, 0.1))
- self._logger.warning(f"Database locked on batch commit, retry {retry_count + 1}/{max_retries} after {delay:.2f}s")
+ delay = cast(
+ float, base_delay * (2**retry_count) + random.uniform(0, 0.1)
+ )
+ self._logger.warning(
+ f"Database locked on batch commit, retry {retry_count + 1}/{max_retries} after {delay:.2f}s"
+ )
await asyncio.sleep(delay)
retry_count += 1
else:
- self._logger.error(f"Failed to commit batch after {retry_count + 1} attempts: {e}")
+ self._logger.error(
+ f"Failed to commit batch after {retry_count + 1} attempts: {e}"
+ )
raise
diff --git a/src/exo/shared/db/sqlite/event_log_manager.py b/src/exo/shared/db/sqlite/event_log_manager.py
index bf09c44c..9a1aa1d9 100644
--- a/src/exo/shared/db/sqlite/event_log_manager.py
+++ b/src/exo/shared/db/sqlite/event_log_manager.py
@@ -13,20 +13,20 @@ class EventLogManager:
"""
Manages both worker and global event log connectors.
Used by both master and worker processes with different access patterns:
-
+
- Worker: writes to worker_events, tails global_events
- Master (elected): writes to global_events, tails global_events
- Master (replica): writes to worker_events, tails global_events
"""
-
+
def __init__(self, config: EventLogConfig, logger: Logger):
self._config = config
self._logger = logger
self._connectors: Dict[EventLogType, AsyncSQLiteEventStorage] = {}
-
+
# Ensure base directory exists
EXO_HOME.mkdir(parents=True, exist_ok=True)
-
+
# TODO: This seems like it's a pattern to avoid an async __init__ function. But as we know, there's a better pattern for this - using a create() function, like in runner_supervisor.
async def initialize(self, max_retries: int = 3) -> None:
"""Initialize both connectors with retry logic - call this during startup"""
@@ -34,7 +34,7 @@ class EventLogManager:
for log_type in [EventLogType.WORKER_EVENTS, EventLogType.GLOBAL_EVENTS]:
retry_count: int = 0
last_error: Optional[Exception] = None
-
+
while retry_count < max_retries:
try:
await self.get_connector(log_type)
@@ -43,26 +43,36 @@ class EventLogManager:
last_error = e
if "database is locked" in str(e) and retry_count < max_retries - 1:
retry_count += 1
- delay = cast(float, 0.5 * (2 ** retry_count))
- self._logger.warning(f"Database locked while initializing {log_type.value}, retry {retry_count}/{max_retries} after {delay}s")
+ delay = cast(float, 0.5 * (2**retry_count))
+ self._logger.warning(
+ f"Database locked while initializing {log_type.value}, retry {retry_count}/{max_retries} after {delay}s"
+ )
await asyncio.sleep(delay)
else:
- self._logger.error(f"Failed to initialize {log_type.value} after {retry_count + 1} attempts: {e}")
- raise RuntimeError(f"Could not initialize {log_type.value} database after {retry_count + 1} attempts") from e
+ self._logger.error(
+ f"Failed to initialize {log_type.value} after {retry_count + 1} attempts: {e}"
+ )
+ raise RuntimeError(
+ f"Could not initialize {log_type.value} database after {retry_count + 1} attempts"
+ ) from e
except Exception as e:
- self._logger.error(f"Unexpected error initializing {log_type.value}: {e}")
+ self._logger.error(
+ f"Unexpected error initializing {log_type.value}: {e}"
+ )
raise
-
+
if retry_count >= max_retries and last_error:
- raise RuntimeError(f"Could not initialize {log_type.value} database after {max_retries} attempts") from last_error
-
+ raise RuntimeError(
+ f"Could not initialize {log_type.value} database after {max_retries} attempts"
+ ) from last_error
+
self._logger.info("Initialized all event log connectors")
-
+
async def get_connector(self, log_type: EventLogType) -> AsyncSQLiteEventStorage:
"""Get or create a connector for the specified log type"""
if log_type not in self._connectors:
db_path = self._config.get_db_path(log_type)
-
+
try:
connector = AsyncSQLiteEventStorage(
db_path=db_path,
@@ -70,37 +80,43 @@ class EventLogManager:
batch_timeout_ms=self._config.batch_timeout_ms,
debounce_ms=self._config.debounce_ms,
max_age_ms=self._config.max_age_ms,
- logger=self._logger
+ logger=self._logger,
)
-
+
# Start the connector (creates tables if needed)
await connector.start()
-
+
self._connectors[log_type] = connector
- self._logger.info(f"Initialized {log_type.value} connector at {db_path}")
+ self._logger.info(
+ f"Initialized {log_type.value} connector at {db_path}"
+ )
except Exception as e:
self._logger.error(f"Failed to create {log_type.value} connector: {e}")
raise
-
+
return self._connectors[log_type]
-
+
@property
def worker_events(self) -> AsyncSQLiteEventStorage:
"""Access worker events log (must call initialize() first)"""
if EventLogType.WORKER_EVENTS not in self._connectors:
- raise RuntimeError("Event log manager not initialized. Call initialize() first.")
+ raise RuntimeError(
+ "Event log manager not initialized. Call initialize() first."
+ )
return self._connectors[EventLogType.WORKER_EVENTS]
-
+
@property
def global_events(self) -> AsyncSQLiteEventStorage:
"""Access global events log (must call initialize() first)"""
if EventLogType.GLOBAL_EVENTS not in self._connectors:
- raise RuntimeError("Event log manager not initialized. Call initialize() first.")
+ raise RuntimeError(
+ "Event log manager not initialized. Call initialize() first."
+ )
return self._connectors[EventLogType.GLOBAL_EVENTS]
-
+
async def close_all(self) -> None:
"""Close all open connectors"""
for log_type, connector in self._connectors.items():
await connector.close()
self._logger.info(f"Closed {log_type.value} connector")
- self._connectors.clear()
\ No newline at end of file
+ self._connectors.clear()
diff --git a/src/exo/shared/db/sqlite/types.py b/src/exo/shared/db/sqlite/types.py
index 3a1cf48e..5fc0f582 100644
--- a/src/exo/shared/db/sqlite/types.py
+++ b/src/exo/shared/db/sqlite/types.py
@@ -11,11 +11,12 @@ from exo.shared.types.events.components import EventFromEventLog
class StoredEvent(SQLModel, table=True):
"""SQLite representation of an event in the event log.
-
+
The rowid serves as the global sequence number (idx_in_log) for ordering.
"""
+
__tablename__ = "events" # type: ignore[assignment]
-
+
# SQLite's rowid as primary key - we alias it but don't actually use it in queries
rowid: int | None = Field(default=None, primary_key=True, alias="rowid")
origin: str = Field(index=True)
@@ -23,39 +24,33 @@ class StoredEvent(SQLModel, table=True):
event_id: str = Field(index=True)
event_data: dict[str, Any] = Field(sa_column=Column(JSON))
created_at: datetime = Field(
- default_factory=lambda: datetime.now(timezone.utc),
- sa_column=Column(DateTime, index=True)
- )
-
- __table_args__ = (
- Index("idx_events_origin_created", "origin", "created_at"),
+ default_factory=lambda: datetime.now(timezone.utc),
+ sa_column=Column(DateTime, index=True),
)
+ __table_args__ = (Index("idx_events_origin_created", "origin", "created_at"),)
+
+
class EventStorageProtocol(Protocol):
"""Protocol for event storage implementations."""
-
- async def append_events(
- self,
- events: Sequence[Event],
- origin: NodeId
- ) -> None:
+
+ async def append_events(self, events: Sequence[Event], origin: NodeId) -> None:
"""Append events to the log (fire-and-forget).
-
+
Events are queued for batched writing and assigned idx_in_log
when committed to storage.
"""
...
-
+
async def get_events_since(
- self,
- last_idx: int
+ self, last_idx: int
) -> Sequence[EventFromEventLog[Event]]:
"""Retrieve events after a specific index.
-
+
Returns events in idx_in_log order.
"""
...
-
+
async def close(self) -> None:
"""Close the storage connection and cleanup resources."""
- ...
\ No newline at end of file
+ ...
diff --git a/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py
index a61d2ecd..1bf4822e 100644
--- a/src/exo/shared/models/model_cards.py
+++ b/src/exo/shared/models/model_cards.py
@@ -15,230 +15,221 @@ class ModelCard(BaseModel):
MODEL_CARDS: dict[str, ModelCard] = {
- # deepseek v3
- "deepseek-v3-0324:4bit": ModelCard(
- short_id="deepseek-v3-0324:4bit",
- model_id="mlx-community/DeepSeek-V3-0324-4bit",
- name="DeepSeek V3 0324 (4-bit)",
- description="""DeepSeek V3 is a large language model trained on the DeepSeek V3 dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id="mlx-community/DeepSeek-V3-0324-4bit",
- pretty_name="DeepSeek V3 0324 (4-bit)",
- storage_size_kilobytes=409706307,
- n_layers=61,
- ),
- ),
- "deepseek-v3-0324": ModelCard(
- short_id="deepseek-v3-0324",
- model_id="mlx-community/DeepSeek-v3-0324-8bit",
- name="DeepSeek V3 0324 (8-bit)",
- description="""DeepSeek V3 is a large language model trained on the DeepSeek V3 dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id="mlx-community/DeepSeek-v3-0324-8bit",
- pretty_name="DeepSeek V3 0324 (8-bit)",
- storage_size_kilobytes=754706307,
- n_layers=61,
- ),
- ),
-
- # deepseek r1
- "deepseek-r1-0528:4bit": ModelCard(
- short_id="deepseek-r1-0528:4bit",
- model_id="mlx-community/DeepSeek-R1-0528-4bit",
- name="DeepSeek-R1-0528 (4-bit)",
- description="""DeepSeek R1 is a large language model trained on the DeepSeek R1 dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id="mlx-community/DeepSeek-R1-0528-4bit",
- pretty_name="DeepSeek R1 671B (4-bit)",
- storage_size_kilobytes=409706307,
- n_layers=61,
- ),
- ),
- "deepseek-r1-0528": ModelCard(
- short_id="deepseek-r1-0528",
- model_id="mlx-community/DeepSeek-R1-0528-8bit",
- name="DeepSeek-R1-0528 (8-bit)",
- description="""DeepSeek R1 is a large language model trained on the DeepSeek R1 dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id="mlx-community/DeepSeek-R1-0528-8bit",
- pretty_name="DeepSeek R1 671B (8-bit)",
- storage_size_kilobytes=754998771712//1024,
- n_layers=61,
- ),
- ),
-
-
- # llama-3.1
- "llama-3.1-8b": ModelCard(
- short_id="llama-3.1-8b",
- model_id="mlx-community/Meta-Llama-3.1-8B-Instruct-4bit",
- name="Llama 3.1 8B",
- description="""Llama 3.1 is a large language model trained on the Llama 3.1 dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id="mlx-community/Meta-Llama-3.1-8B-Instruct-4bit",
- pretty_name="Llama 3.1 8B",
- storage_size_kilobytes=4411528,
- n_layers=32,
- ),
- ),
- "llama-3.1-70b": ModelCard(
- short_id="llama-3.1-70b",
- model_id="mlx-community/Meta-Llama-3.1-70B-Instruct-4bit",
- name="Llama 3.1 70B",
- description="""Llama 3.1 is a large language model trained on the Llama 3.1 dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id="mlx-community/Meta-Llama-3.1-70B-Instruct-4bit",
- pretty_name="Llama 3.1 70B",
- storage_size_kilobytes=38758160,
- n_layers=80,
- ),
- ),
-
- # llama-3.2
- "llama-3.2-1b": ModelCard(
- short_id="llama-3.2-1b",
- model_id="mlx-community/Llama-3.2-1B-Instruct-4bit",
- name="Llama 3.2 1B",
- description="""Llama 3.2 is a large language model trained on the Llama 3.2 dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id="mlx-community/Llama-3.2-1B-Instruct-4bit",
- pretty_name="Llama 3.2 1B",
- storage_size_kilobytes=678948,
- n_layers=16,
- ),
- ),
- "llama-3.2-3b": ModelCard(
- short_id="llama-3.2-3b",
- model_id="mlx-community/Llama-3.2-3B-Instruct-4bit",
- name="Llama 3.2 3B",
- description="""Llama 3.2 is a large language model trained on the Llama 3.2 dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id="mlx-community/Llama-3.2-3B-Instruct-4bit",
- pretty_name="Llama 3.2 3B",
- storage_size_kilobytes=1765062,
- n_layers=28,
- ),
- ),
-
- # llama-3.3
- "llama-3.3-70b": ModelCard(
- short_id="llama-3.3-70b",
- model_id="mlx-community/Llama-3.3-70B-Instruct-4bit",
- name="Llama 3.3 70B",
- description="""The Meta Llama 3.3 multilingual large language model (LLM) is an instruction tuned generative model in 70B (text in/text out)""",
- tags=[],
- metadata=ModelMetadata(
- model_id="mlx-community/Llama-3.3-70B-Instruct-4bit",
- pretty_name="Llama 3.3 70B",
- storage_size_kilobytes=38758160,
- n_layers=80,
- ),
- ),
-
- # phi-3
- "phi-3-mini": ModelCard(
- short_id="phi-3-mini",
- model_id="mlx-community/Phi-3-mini-128k-instruct-4bit",
- name="Phi 3 Mini 128k",
- description="""Phi 3 Mini is a large language model trained on the Phi 3 Mini dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id="mlx-community/Phi-3-mini-128k-instruct-4bit",
- pretty_name="Phi 3 Mini 128k",
- storage_size_kilobytes=2099262,
- n_layers=32,
- ),
- ),
- "phi-3-mini:128k": ModelCard(
- short_id="phi-3-mini:128k",
- model_id="mlx-community/Phi-3-mini-128k-instruct-4bit",
- name="Phi 3 Mini 128k",
- description="""Phi 3 Mini is a large language model trained on the Phi 3 Mini dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id="mlx-community/Phi-3-mini-128k-instruct-4bit",
- pretty_name="Phi 3 Mini 128k",
- storage_size_kilobytes=2099262,
- n_layers=32,
- ),
- ),
-
- # qwen3
- "qwen3-0.6b": ModelCard(
- short_id="qwen3-0.6b",
- model_id="mlx-community/Qwen3-0.6B-4bit",
- name="Qwen3 0.6B",
- description="""Qwen3 0.6B is a large language model trained on the Qwen3 0.6B dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id="mlx-community/Qwen3-0.6B-4bit",
- pretty_name="Qwen3 0.6B",
- storage_size_kilobytes=327512,
- n_layers=28,
- ),
- ),
- "qwen3-30b": ModelCard(
- short_id="qwen3-30b",
- model_id="mlx-community/Qwen3-30B-A3B-4bit",
- name="Qwen3 30B (Active 3B)",
- description="""Qwen3 30B is a large language model trained on the Qwen3 30B dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id="mlx-community/Qwen3-30B-A3B-4bit",
- pretty_name="Qwen3 30B (Active 3B)",
- storage_size_kilobytes=16772092,
- n_layers=48,
- ),
- ),
-
- # granite
- "granite-3.3-2b": ModelCard(
- short_id="granite-3.3-2b",
- model_id="mlx-community/granite-3.3-2b-instruct-fp16",
- name="Granite 3.3 2B",
- description="""Granite-3.3-2B-Instruct is a 2-billion parameter 128K context length language model fine-tuned for improved reasoning and instruction-following capabilities.""",
- tags=[],
- metadata=ModelMetadata(
- model_id="mlx-community/granite-3.3-2b-instruct-fp16",
- pretty_name="Granite 3.3 2B",
- storage_size_kilobytes=4948320,
- n_layers=40,
- ),
- ),
- "granite-3.3-8b": ModelCard(
- short_id="granite-3.3-8b",
- model_id="mlx-community/granite-3.3-8b-instruct-fp16",
- name="Granite 3.3 8B",
- description="""Granite-3.3-8B-Instruct is a 8-billion parameter 128K context length language model fine-tuned for improved reasoning and instruction-following capabilities.""",
- tags=[],
- metadata=ModelMetadata(
- model_id="mlx-community/granite-3.3-8b-instruct-fp16",
- pretty_name="Granite 3.3 8B",
- storage_size_kilobytes=15958720,
- n_layers=40,
- ),
- ),
-
- # smol-lm
- "smol-lm-135m": ModelCard(
- short_id="smol-lm-135m",
- model_id="mlx-community/SmolLM-135M-4bit",
- name="Smol LM 135M",
- description="""SmolLM is a series of state-of-the-art small language models available in three sizes: 135M, 360M, and 1.7B parameters. """,
- tags=[],
- metadata=ModelMetadata(
- model_id="mlx-community/SmolLM-135M-4bit",
- pretty_name="Smol LM 135M",
- storage_size_kilobytes=73940,
- n_layers=30,
- ),
- ),
+ # deepseek v3
+ "deepseek-v3-0324:4bit": ModelCard(
+ short_id="deepseek-v3-0324:4bit",
+ model_id="mlx-community/DeepSeek-V3-0324-4bit",
+ name="DeepSeek V3 0324 (4-bit)",
+ description="""DeepSeek V3 is a large language model trained on the DeepSeek V3 dataset.""",
+ tags=[],
+ metadata=ModelMetadata(
+ model_id="mlx-community/DeepSeek-V3-0324-4bit",
+ pretty_name="DeepSeek V3 0324 (4-bit)",
+ storage_size_kilobytes=409706307,
+ n_layers=61,
+ ),
+ ),
+ "deepseek-v3-0324": ModelCard(
+ short_id="deepseek-v3-0324",
+ model_id="mlx-community/DeepSeek-v3-0324-8bit",
+ name="DeepSeek V3 0324 (8-bit)",
+ description="""DeepSeek V3 is a large language model trained on the DeepSeek V3 dataset.""",
+ tags=[],
+ metadata=ModelMetadata(
+ model_id="mlx-community/DeepSeek-v3-0324-8bit",
+ pretty_name="DeepSeek V3 0324 (8-bit)",
+ storage_size_kilobytes=754706307,
+ n_layers=61,
+ ),
+ ),
+ # deepseek r1
+ "deepseek-r1-0528:4bit": ModelCard(
+ short_id="deepseek-r1-0528:4bit",
+ model_id="mlx-community/DeepSeek-R1-0528-4bit",
+ name="DeepSeek-R1-0528 (4-bit)",
+ description="""DeepSeek R1 is a large language model trained on the DeepSeek R1 dataset.""",
+ tags=[],
+ metadata=ModelMetadata(
+ model_id="mlx-community/DeepSeek-R1-0528-4bit",
+ pretty_name="DeepSeek R1 671B (4-bit)",
+ storage_size_kilobytes=409706307,
+ n_layers=61,
+ ),
+ ),
+ "deepseek-r1-0528": ModelCard(
+ short_id="deepseek-r1-0528",
+ model_id="mlx-community/DeepSeek-R1-0528-8bit",
+ name="DeepSeek-R1-0528 (8-bit)",
+ description="""DeepSeek R1 is a large language model trained on the DeepSeek R1 dataset.""",
+ tags=[],
+ metadata=ModelMetadata(
+ model_id="mlx-community/DeepSeek-R1-0528-8bit",
+ pretty_name="DeepSeek R1 671B (8-bit)",
+ storage_size_kilobytes=754998771712 // 1024,
+ n_layers=61,
+ ),
+ ),
+ # llama-3.1
+ "llama-3.1-8b": ModelCard(
+ short_id="llama-3.1-8b",
+ model_id="mlx-community/Meta-Llama-3.1-8B-Instruct-4bit",
+ name="Llama 3.1 8B",
+ description="""Llama 3.1 is a large language model trained on the Llama 3.1 dataset.""",
+ tags=[],
+ metadata=ModelMetadata(
+ model_id="mlx-community/Meta-Llama-3.1-8B-Instruct-4bit",
+ pretty_name="Llama 3.1 8B",
+ storage_size_kilobytes=4411528,
+ n_layers=32,
+ ),
+ ),
+ "llama-3.1-70b": ModelCard(
+ short_id="llama-3.1-70b",
+ model_id="mlx-community/Meta-Llama-3.1-70B-Instruct-4bit",
+ name="Llama 3.1 70B",
+ description="""Llama 3.1 is a large language model trained on the Llama 3.1 dataset.""",
+ tags=[],
+ metadata=ModelMetadata(
+ model_id="mlx-community/Meta-Llama-3.1-70B-Instruct-4bit",
+ pretty_name="Llama 3.1 70B",
+ storage_size_kilobytes=38758160,
+ n_layers=80,
+ ),
+ ),
+ # llama-3.2
+ "llama-3.2-1b": ModelCard(
+ short_id="llama-3.2-1b",
+ model_id="mlx-community/Llama-3.2-1B-Instruct-4bit",
+ name="Llama 3.2 1B",
+ description="""Llama 3.2 is a large language model trained on the Llama 3.2 dataset.""",
+ tags=[],
+ metadata=ModelMetadata(
+ model_id="mlx-community/Llama-3.2-1B-Instruct-4bit",
+ pretty_name="Llama 3.2 1B",
+ storage_size_kilobytes=678948,
+ n_layers=16,
+ ),
+ ),
+ "llama-3.2-3b": ModelCard(
+ short_id="llama-3.2-3b",
+ model_id="mlx-community/Llama-3.2-3B-Instruct-4bit",
+ name="Llama 3.2 3B",
+ description="""Llama 3.2 is a large language model trained on the Llama 3.2 dataset.""",
+ tags=[],
+ metadata=ModelMetadata(
+ model_id="mlx-community/Llama-3.2-3B-Instruct-4bit",
+ pretty_name="Llama 3.2 3B",
+ storage_size_kilobytes=1765062,
+ n_layers=28,
+ ),
+ ),
+ # llama-3.3
+ "llama-3.3-70b": ModelCard(
+ short_id="llama-3.3-70b",
+ model_id="mlx-community/Llama-3.3-70B-Instruct-4bit",
+ name="Llama 3.3 70B",
+ description="""The Meta Llama 3.3 multilingual large language model (LLM) is an instruction tuned generative model in 70B (text in/text out)""",
+ tags=[],
+ metadata=ModelMetadata(
+ model_id="mlx-community/Llama-3.3-70B-Instruct-4bit",
+ pretty_name="Llama 3.3 70B",
+ storage_size_kilobytes=38758160,
+ n_layers=80,
+ ),
+ ),
+ # phi-3
+ "phi-3-mini": ModelCard(
+ short_id="phi-3-mini",
+ model_id="mlx-community/Phi-3-mini-128k-instruct-4bit",
+ name="Phi 3 Mini 128k",
+ description="""Phi 3 Mini is a large language model trained on the Phi 3 Mini dataset.""",
+ tags=[],
+ metadata=ModelMetadata(
+ model_id="mlx-community/Phi-3-mini-128k-instruct-4bit",
+ pretty_name="Phi 3 Mini 128k",
+ storage_size_kilobytes=2099262,
+ n_layers=32,
+ ),
+ ),
+ "phi-3-mini:128k": ModelCard(
+ short_id="phi-3-mini:128k",
+ model_id="mlx-community/Phi-3-mini-128k-instruct-4bit",
+ name="Phi 3 Mini 128k",
+ description="""Phi 3 Mini is a large language model trained on the Phi 3 Mini dataset.""",
+ tags=[],
+ metadata=ModelMetadata(
+ model_id="mlx-community/Phi-3-mini-128k-instruct-4bit",
+ pretty_name="Phi 3 Mini 128k",
+ storage_size_kilobytes=2099262,
+ n_layers=32,
+ ),
+ ),
+ # qwen3
+ "qwen3-0.6b": ModelCard(
+ short_id="qwen3-0.6b",
+ model_id="mlx-community/Qwen3-0.6B-4bit",
+ name="Qwen3 0.6B",
+ description="""Qwen3 0.6B is a large language model trained on the Qwen3 0.6B dataset.""",
+ tags=[],
+ metadata=ModelMetadata(
+ model_id="mlx-community/Qwen3-0.6B-4bit",
+ pretty_name="Qwen3 0.6B",
+ storage_size_kilobytes=327512,
+ n_layers=28,
+ ),
+ ),
+ "qwen3-30b": ModelCard(
+ short_id="qwen3-30b",
+ model_id="mlx-community/Qwen3-30B-A3B-4bit",
+ name="Qwen3 30B (Active 3B)",
+ description="""Qwen3 30B is a large language model trained on the Qwen3 30B dataset.""",
+ tags=[],
+ metadata=ModelMetadata(
+ model_id="mlx-community/Qwen3-30B-A3B-4bit",
+ pretty_name="Qwen3 30B (Active 3B)",
+ storage_size_kilobytes=16772092,
+ n_layers=48,
+ ),
+ ),
+ # granite
+ "granite-3.3-2b": ModelCard(
+ short_id="granite-3.3-2b",
+ model_id="mlx-community/granite-3.3-2b-instruct-fp16",
+ name="Granite 3.3 2B",
+ description="""Granite-3.3-2B-Instruct is a 2-billion parameter 128K context length language model fine-tuned for improved reasoning and instruction-following capabilities.""",
+ tags=[],
+ metadata=ModelMetadata(
+ model_id="mlx-community/granite-3.3-2b-instruct-fp16",
+ pretty_name="Granite 3.3 2B",
+ storage_size_kilobytes=4948320,
+ n_layers=40,
+ ),
+ ),
+ "granite-3.3-8b": ModelCard(
+ short_id="granite-3.3-8b",
+ model_id="mlx-community/granite-3.3-8b-instruct-fp16",
+ name="Granite 3.3 8B",
+ description="""Granite-3.3-8B-Instruct is a 8-billion parameter 128K context length language model fine-tuned for improved reasoning and instruction-following capabilities.""",
+ tags=[],
+ metadata=ModelMetadata(
+ model_id="mlx-community/granite-3.3-8b-instruct-fp16",
+ pretty_name="Granite 3.3 8B",
+ storage_size_kilobytes=15958720,
+ n_layers=40,
+ ),
+ ),
+ # smol-lm
+ "smol-lm-135m": ModelCard(
+ short_id="smol-lm-135m",
+ model_id="mlx-community/SmolLM-135M-4bit",
+ name="Smol LM 135M",
+ description="""SmolLM is a series of state-of-the-art small language models available in three sizes: 135M, 360M, and 1.7B parameters. """,
+ tags=[],
+ metadata=ModelMetadata(
+ model_id="mlx-community/SmolLM-135M-4bit",
+ pretty_name="Smol LM 135M",
+ storage_size_kilobytes=73940,
+ n_layers=30,
+ ),
+ ),
}
diff --git a/src/exo/shared/models/model_meta.py b/src/exo/shared/models/model_meta.py
index 57532053..31260eae 100644
--- a/src/exo/shared/models/model_meta.py
+++ b/src/exo/shared/models/model_meta.py
@@ -21,7 +21,9 @@ class ConfigData(BaseModel):
num_layers: Optional[Annotated[int, Field(ge=0)]] = None
n_layer: Optional[Annotated[int, Field(ge=0)]] = None
n_layers: Optional[Annotated[int, Field(ge=0)]] = None # Sometimes used
- num_decoder_layers: Optional[Annotated[int, Field(ge=0)]] = None # Transformer models
+ num_decoder_layers: Optional[Annotated[int, Field(ge=0)]] = (
+ None # Transformer models
+ )
decoder_layers: Optional[Annotated[int, Field(ge=0)]] = None # Some architectures
@property
@@ -40,22 +42,42 @@ class ConfigData(BaseModel):
if layer_count is not None:
return layer_count
- raise ValueError(f"No layer count found in config.json: {self.model_dump_json()}")
+ raise ValueError(
+ f"No layer count found in config.json: {self.model_dump_json()}"
+ )
+
async def get_config_data(model_id: str) -> ConfigData:
"""Downloads and parses config.json for a model."""
- target_dir = (await ensure_models_dir())/str(model_id).replace("/", "--")
+ target_dir = (await ensure_models_dir()) / str(model_id).replace("/", "--")
await aios.makedirs(target_dir, exist_ok=True)
- config_path = await download_file_with_retry(model_id, "main", "config.json", target_dir, lambda curr_bytes, total_bytes: print(f"Downloading config.json for {model_id}: {curr_bytes}/{total_bytes}"))
- async with aiofiles.open(config_path, 'r') as f:
+ config_path = await download_file_with_retry(
+ model_id,
+ "main",
+ "config.json",
+ target_dir,
+ lambda curr_bytes, total_bytes: print(
+ f"Downloading config.json for {model_id}: {curr_bytes}/{total_bytes}"
+ ),
+ )
+ async with aiofiles.open(config_path, "r") as f:
return ConfigData.model_validate_json(await f.read())
+
async def get_safetensors_size(model_id: str) -> int:
"""Gets model size from safetensors index or falls back to HF API."""
- target_dir = (await ensure_models_dir())/str(model_id).replace("/", "--")
+ target_dir = (await ensure_models_dir()) / str(model_id).replace("/", "--")
await aios.makedirs(target_dir, exist_ok=True)
- index_path = await download_file_with_retry(model_id, "main", "model.safetensors.index.json", target_dir, lambda curr_bytes, total_bytes: print(f"Downloading model.safetensors.index.json for {model_id}: {curr_bytes}/{total_bytes}"))
- async with aiofiles.open(index_path, 'r') as f:
+ index_path = await download_file_with_retry(
+ model_id,
+ "main",
+ "model.safetensors.index.json",
+ target_dir,
+ lambda curr_bytes, total_bytes: print(
+ f"Downloading model.safetensors.index.json for {model_id}: {curr_bytes}/{total_bytes}"
+ ),
+ )
+ async with aiofiles.open(index_path, "r") as f:
index_data = ModelSafetensorsIndex.model_validate_json(await f.read())
metadata = index_data.metadata
@@ -67,7 +89,10 @@ async def get_safetensors_size(model_id: str) -> int:
raise ValueError(f"No safetensors info found for {model_id}")
return info.safetensors.total
+
_model_meta_cache: Dict[str, ModelMetadata] = {}
+
+
async def get_model_meta(model_id: str) -> ModelMetadata:
if model_id in _model_meta_cache:
return _model_meta_cache[model_id]
@@ -75,6 +100,7 @@ async def get_model_meta(model_id: str) -> ModelMetadata:
_model_meta_cache[model_id] = model_meta
return model_meta
+
async def _get_model_meta(model_id: str) -> ModelMetadata:
"""Fetches storage size and number of layers for a Hugging Face model, returns Pydantic ModelMeta."""
config_data = await get_config_data(model_id)
diff --git a/src/exo/shared/tests/__init__.py b/src/exo/shared/tests/__init__.py
index e5374d95..09c36e8f 100644
--- a/src/exo/shared/tests/__init__.py
+++ b/src/exo/shared/tests/__init__.py
@@ -1 +1 @@
-# Test package for shared utilities
\ No newline at end of file
+# Test package for shared utilities
diff --git a/src/exo/shared/tests/test_node_id_persistence.py b/src/exo/shared/tests/test_node_id_persistence.py
index 1f41cf99..552311e7 100644
--- a/src/exo/shared/tests/test_node_id_persistence.py
+++ b/src/exo/shared/tests/test_node_id_persistence.py
@@ -19,7 +19,9 @@ from exo.shared.utils import get_node_id_keypair
NUM_CONCURRENT_PROCS = 10
-def _get_keypair_concurrent_subprocess_task(pid: int, sem: SemaphoreT, ev: EventT, queue: QueueT[bytes]) -> None:
+def _get_keypair_concurrent_subprocess_task(
+ pid: int, sem: SemaphoreT, ev: EventT, queue: QueueT[bytes]
+) -> None:
try:
# synchronise with parent process
logging.info(msg=f"SUBPROCESS {pid}: Started")
@@ -45,8 +47,9 @@ def _get_keypair_concurrent(num_procs: int) -> bytes:
logging.info(msg=f"PARENT: Starting {num_procs} subprocesses")
ps: list[BaseProcess] = []
for i in range(num_procs):
- p = multiprocessing.get_context("fork").Process(target=_get_keypair_concurrent_subprocess_task,
- args=(i + 1, sem, ev, queue))
+ p = multiprocessing.get_context("fork").Process(
+ target=_get_keypair_concurrent_subprocess_task, args=(i + 1, sem, ev, queue)
+ )
ps.append(p)
p.start()
for _ in range(num_procs):
diff --git a/src/exo/shared/tests/test_sqlite_connector.py b/src/exo/shared/tests/test_sqlite_connector.py
index 30979e5c..8917e9ce 100644
--- a/src/exo/shared/tests/test_sqlite_connector.py
+++ b/src/exo/shared/tests/test_sqlite_connector.py
@@ -45,57 +45,83 @@ class TestAsyncSQLiteEventStorage:
async def test_initialization_creates_tables(self, temp_db_path: Path) -> None:
"""Test that database initialization creates the events table."""
default_config = EventLogConfig()
- storage = AsyncSQLiteEventStorage(db_path=temp_db_path, batch_size=default_config.batch_size, batch_timeout_ms=default_config.batch_timeout_ms, debounce_ms=default_config.debounce_ms, max_age_ms=default_config.max_age_ms)
+ storage = AsyncSQLiteEventStorage(
+ db_path=temp_db_path,
+ batch_size=default_config.batch_size,
+ batch_timeout_ms=default_config.batch_timeout_ms,
+ debounce_ms=default_config.debounce_ms,
+ max_age_ms=default_config.max_age_ms,
+ )
await storage.start()
-
+
# Verify table exists by querying directly
assert storage._engine is not None
async with AsyncSession(storage._engine) as session:
- result = await session.execute(text("SELECT name FROM sqlite_master WHERE type='table' AND name='events'"))
+ result = await session.execute(
+ text(
+ "SELECT name FROM sqlite_master WHERE type='table' AND name='events'"
+ )
+ )
tables = result.fetchall()
assert len(tables) == 1
assert tables[0][0] == "events"
-
+
await storage.close()
@pytest.mark.asyncio
async def test_start_twice_raises_error(self, temp_db_path: Path) -> None:
"""Test that starting storage twice raises an error."""
default_config = EventLogConfig()
- storage = AsyncSQLiteEventStorage(db_path=temp_db_path, batch_size=default_config.batch_size, batch_timeout_ms=default_config.batch_timeout_ms, debounce_ms=default_config.debounce_ms, max_age_ms=default_config.max_age_ms)
+ storage = AsyncSQLiteEventStorage(
+ db_path=temp_db_path,
+ batch_size=default_config.batch_size,
+ batch_timeout_ms=default_config.batch_timeout_ms,
+ debounce_ms=default_config.debounce_ms,
+ max_age_ms=default_config.max_age_ms,
+ )
await storage.start()
-
+
with pytest.raises(RuntimeError, match="Storage already started"):
await storage.start()
-
+
await storage.close()
@pytest.mark.asyncio
- async def test_direct_database_operations(self, temp_db_path: Path, sample_node_id: NodeId) -> None:
+ async def test_direct_database_operations(
+ self, temp_db_path: Path, sample_node_id: NodeId
+ ) -> None:
"""Test direct database operations without event parsing."""
default_config = EventLogConfig()
- storage = AsyncSQLiteEventStorage(db_path=temp_db_path, batch_size=default_config.batch_size, batch_timeout_ms=default_config.batch_timeout_ms, debounce_ms=default_config.debounce_ms, max_age_ms=default_config.max_age_ms)
+ storage = AsyncSQLiteEventStorage(
+ db_path=temp_db_path,
+ batch_size=default_config.batch_size,
+ batch_timeout_ms=default_config.batch_timeout_ms,
+ debounce_ms=default_config.debounce_ms,
+ max_age_ms=default_config.max_age_ms,
+ )
await storage.start()
-
+
# Insert test data directly
test_data = {
"event_type": "test_event",
"test_field": "test_value",
- "number": 42
+ "number": 42,
}
-
+
async with AsyncSession(storage._engine) as session:
await session.execute(
- text("INSERT INTO events (origin, event_type, event_id, event_data) VALUES (:origin, :event_type, :event_id, :event_data)"),
+ text(
+ "INSERT INTO events (origin, event_type, event_id, event_data) VALUES (:origin, :event_type, :event_id, :event_data)"
+ ),
{
"origin": sample_node_id,
"event_type": "test_event",
"event_id": str(uuid4()),
- "event_data": json.dumps(test_data)
- }
+ "event_data": json.dumps(test_data),
+ },
)
await session.commit()
-
+
# Query data back
assert storage._engine is not None
async with AsyncSession(storage._engine) as session:
@@ -103,44 +129,54 @@ class TestAsyncSQLiteEventStorage:
text("SELECT rowid, origin, event_data FROM events ORDER BY rowid")
)
rows = result.fetchall()
-
+
assert len(rows) == 1
assert rows[0][0] == 1 # rowid
assert rows[0][1] == sample_node_id # origin
raw_json = cast(str, rows[0][2])
retrieved_data = _load_json_data(raw_json)
assert retrieved_data == test_data
-
+
await storage.close()
@pytest.mark.asyncio
- async def test_rowid_auto_increment(self, temp_db_path: Path, sample_node_id: NodeId) -> None:
+ async def test_rowid_auto_increment(
+ self, temp_db_path: Path, sample_node_id: NodeId
+ ) -> None:
"""Test that rowid auto-increments correctly."""
default_config = EventLogConfig()
- storage = AsyncSQLiteEventStorage(db_path=temp_db_path, batch_size=default_config.batch_size, batch_timeout_ms=default_config.batch_timeout_ms, debounce_ms=default_config.debounce_ms, max_age_ms=default_config.max_age_ms)
+ storage = AsyncSQLiteEventStorage(
+ db_path=temp_db_path,
+ batch_size=default_config.batch_size,
+ batch_timeout_ms=default_config.batch_timeout_ms,
+ debounce_ms=default_config.debounce_ms,
+ max_age_ms=default_config.max_age_ms,
+ )
await storage.start()
-
+
# Insert multiple records
test_records = [
{"event_type": "test_event_1", "data": "first"},
{"event_type": "test_event_2", "data": "second"},
- {"event_type": "test_event_3", "data": "third"}
+ {"event_type": "test_event_3", "data": "third"},
]
-
+
assert storage._engine is not None
async with AsyncSession(storage._engine) as session:
for record in test_records:
await session.execute(
- text("INSERT INTO events (origin, event_type, event_id, event_data) VALUES (:origin, :event_type, :event_id, :event_data)"),
+ text(
+ "INSERT INTO events (origin, event_type, event_id, event_data) VALUES (:origin, :event_type, :event_id, :event_data)"
+ ),
{
"origin": sample_node_id,
"event_type": record["event_type"],
"event_id": str(uuid4()),
- "event_data": json.dumps(record)
- }
+ "event_data": json.dumps(record),
+ },
)
await session.commit()
-
+
# Query back and verify rowid sequence
assert storage._engine is not None
async with AsyncSession(storage._engine) as session:
@@ -148,81 +184,116 @@ class TestAsyncSQLiteEventStorage:
text("SELECT rowid, event_data FROM events ORDER BY rowid")
)
rows = result.fetchall()
-
+
assert len(rows) == 3
for i, row in enumerate(rows):
assert row[0] == i + 1 # rowid starts at 1
raw_json = cast(str, row[1])
retrieved_data = _load_json_data(raw_json)
assert retrieved_data == test_records[i]
-
- await storage.close()
-
+ await storage.close()
@pytest.mark.asyncio
- async def test_get_last_idx(self, temp_db_path: Path, sample_node_id: NodeId) -> None:
+ async def test_get_last_idx(
+ self, temp_db_path: Path, sample_node_id: NodeId
+ ) -> None:
"""Test that rowid returns correctly from db."""
default_config = EventLogConfig()
- storage = AsyncSQLiteEventStorage(db_path=temp_db_path, batch_size=default_config.batch_size, batch_timeout_ms=default_config.batch_timeout_ms, debounce_ms=default_config.debounce_ms, max_age_ms=default_config.max_age_ms)
+ storage = AsyncSQLiteEventStorage(
+ db_path=temp_db_path,
+ batch_size=default_config.batch_size,
+ batch_timeout_ms=default_config.batch_timeout_ms,
+ debounce_ms=default_config.debounce_ms,
+ max_age_ms=default_config.max_age_ms,
+ )
await storage.start()
-
+
# Insert multiple records
test_records = [
{"event_type": "test_event_1", "data": "first"},
{"event_type": "test_event_2", "data": "second"},
- {"event_type": "test_event_3", "data": "third"}
+ {"event_type": "test_event_3", "data": "third"},
]
-
+
assert storage._engine is not None
async with AsyncSession(storage._engine) as session:
for record in test_records:
await session.execute(
- text("INSERT INTO events (origin, event_type, event_id, event_data) VALUES (:origin, :event_type, :event_id, :event_data)"),
+ text(
+ "INSERT INTO events (origin, event_type, event_id, event_data) VALUES (:origin, :event_type, :event_id, :event_data)"
+ ),
{
"origin": sample_node_id,
"event_type": record["event_type"],
"event_id": str(uuid4()),
- "event_data": json.dumps(record)
- }
+ "event_data": json.dumps(record),
+ },
)
await session.commit()
-
+
last_idx = await storage.get_last_idx()
assert last_idx == 3
-
+
await storage.close()
@pytest.mark.asyncio
async def test_rowid_with_multiple_origins(self, temp_db_path: Path) -> None:
"""Test rowid sequence across multiple origins."""
default_config = EventLogConfig()
- storage = AsyncSQLiteEventStorage(db_path=temp_db_path, batch_size=default_config.batch_size, batch_timeout_ms=default_config.batch_timeout_ms, debounce_ms=default_config.debounce_ms, max_age_ms=default_config.max_age_ms)
+ storage = AsyncSQLiteEventStorage(
+ db_path=temp_db_path,
+ batch_size=default_config.batch_size,
+ batch_timeout_ms=default_config.batch_timeout_ms,
+ debounce_ms=default_config.debounce_ms,
+ max_age_ms=default_config.max_age_ms,
+ )
await storage.start()
-
+
origin1 = NodeId()
origin2 = NodeId()
-
+
# Insert interleaved records from different origins
assert storage._engine is not None
async with AsyncSession(storage._engine) as session:
# Origin 1 - record 1
await session.execute(
- text("INSERT INTO events (origin, event_type, event_id, event_data) VALUES (:origin, :event_type, :event_id, :event_data)"),
- {"origin": origin1, "event_type": "event_1", "event_id": str(uuid4()), "event_data": json.dumps({"from": "origin1", "seq": 1})}
+ text(
+ "INSERT INTO events (origin, event_type, event_id, event_data) VALUES (:origin, :event_type, :event_id, :event_data)"
+ ),
+ {
+ "origin": origin1,
+ "event_type": "event_1",
+ "event_id": str(uuid4()),
+ "event_data": json.dumps({"from": "origin1", "seq": 1}),
+ },
)
# Origin 2 - record 2
await session.execute(
- text("INSERT INTO events (origin, event_type, event_id, event_data) VALUES (:origin, :event_type, :event_id, :event_data)"),
- {"origin": origin2, "event_type": "event_2", "event_id": str(uuid4()), "event_data": json.dumps({"from": "origin2", "seq": 2})}
+ text(
+ "INSERT INTO events (origin, event_type, event_id, event_data) VALUES (:origin, :event_type, :event_id, :event_data)"
+ ),
+ {
+ "origin": origin2,
+ "event_type": "event_2",
+ "event_id": str(uuid4()),
+ "event_data": json.dumps({"from": "origin2", "seq": 2}),
+ },
)
# Origin 1 - record 3
await session.execute(
- text("INSERT INTO events (origin, event_type, event_id, event_data) VALUES (:origin, :event_type, :event_id, :event_data)"),
- {"origin": origin1, "event_type": "event_3", "event_id": str(uuid4()), "event_data": json.dumps({"from": "origin1", "seq": 3})}
+ text(
+ "INSERT INTO events (origin, event_type, event_id, event_data) VALUES (:origin, :event_type, :event_id, :event_data)"
+ ),
+ {
+ "origin": origin1,
+ "event_type": "event_3",
+ "event_id": str(uuid4()),
+ "event_data": json.dumps({"from": "origin1", "seq": 3}),
+ },
)
await session.commit()
-
+
# Verify sequential rowid regardless of origin
assert storage._engine is not None
async with AsyncSession(storage._engine) as session:
@@ -230,12 +301,12 @@ class TestAsyncSQLiteEventStorage:
text("SELECT rowid, origin, event_data FROM events ORDER BY rowid")
)
rows = result.fetchall()
-
+
assert len(rows) == 3
assert rows[0][0] == 1 # First rowid
assert rows[1][0] == 2 # Second rowid
assert rows[2][0] == 3 # Third rowid
-
+
# Verify data integrity
raw_json1 = cast(str, rows[0][2])
raw_json2 = cast(str, rows[1][2])
@@ -243,80 +314,106 @@ class TestAsyncSQLiteEventStorage:
data1 = _load_json_data(raw_json1)
data2 = _load_json_data(raw_json2)
data3 = _load_json_data(raw_json3)
-
+
assert data1["from"] == "origin1" and data1["seq"] == 1
assert data2["from"] == "origin2" and data2["seq"] == 2
assert data3["from"] == "origin1" and data3["seq"] == 3
-
+
await storage.close()
@pytest.mark.asyncio
- async def test_query_events_since_index(self, temp_db_path: Path, sample_node_id: NodeId) -> None:
+ async def test_query_events_since_index(
+ self, temp_db_path: Path, sample_node_id: NodeId
+ ) -> None:
"""Test querying events after a specific rowid."""
default_config = EventLogConfig()
- storage = AsyncSQLiteEventStorage(db_path=temp_db_path, batch_size=default_config.batch_size, batch_timeout_ms=default_config.batch_timeout_ms, debounce_ms=default_config.debounce_ms, max_age_ms=default_config.max_age_ms)
+ storage = AsyncSQLiteEventStorage(
+ db_path=temp_db_path,
+ batch_size=default_config.batch_size,
+ batch_timeout_ms=default_config.batch_timeout_ms,
+ debounce_ms=default_config.debounce_ms,
+ max_age_ms=default_config.max_age_ms,
+ )
await storage.start()
-
+
# Insert 10 test records
assert storage._engine is not None
async with AsyncSession(storage._engine) as session:
for i in range(10):
await session.execute(
- text("INSERT INTO events (origin, event_type, event_id, event_data) VALUES (:origin, :event_type, :event_id, :event_data)"),
+ text(
+ "INSERT INTO events (origin, event_type, event_id, event_data) VALUES (:origin, :event_type, :event_id, :event_data)"
+ ),
{
"origin": sample_node_id,
"event_type": f"event_{i}",
"event_id": str(uuid4()),
- "event_data": json.dumps({"index": i})
- }
+ "event_data": json.dumps({"index": i}),
+ },
)
await session.commit()
-
+
# Query events after index 5
assert storage._engine is not None
async with AsyncSession(storage._engine) as session:
result = await session.execute(
- text("SELECT rowid, event_data FROM events WHERE rowid > :last_idx ORDER BY rowid"),
- {"last_idx": 5}
+ text(
+ "SELECT rowid, event_data FROM events WHERE rowid > :last_idx ORDER BY rowid"
+ ),
+ {"last_idx": 5},
)
rows = result.fetchall()
-
+
assert len(rows) == 5 # Should get records 6-10
for i, row in enumerate(rows):
assert row[0] == i + 6 # rowid 6, 7, 8, 9, 10
raw_json = cast(str, row[1])
data = _load_json_data(raw_json)
assert data["index"] == i + 5 # index 5, 6, 7, 8, 9
-
+
await storage.close()
@pytest.mark.asyncio
async def test_empty_query(self, temp_db_path: Path) -> None:
"""Test querying when no events exist."""
default_config = EventLogConfig()
- storage = AsyncSQLiteEventStorage(db_path=temp_db_path, batch_size=default_config.batch_size, batch_timeout_ms=default_config.batch_timeout_ms, debounce_ms=default_config.debounce_ms, max_age_ms=default_config.max_age_ms)
+ storage = AsyncSQLiteEventStorage(
+ db_path=temp_db_path,
+ batch_size=default_config.batch_size,
+ batch_timeout_ms=default_config.batch_timeout_ms,
+ debounce_ms=default_config.debounce_ms,
+ max_age_ms=default_config.max_age_ms,
+ )
await storage.start()
-
+
assert storage._engine is not None
async with AsyncSession(storage._engine) as session:
result = await session.execute(
- text("SELECT rowid, origin, event_data FROM events WHERE rowid > :last_idx ORDER BY rowid"),
- {"last_idx": 0}
+ text(
+ "SELECT rowid, origin, event_data FROM events WHERE rowid > :last_idx ORDER BY rowid"
+ ),
+ {"last_idx": 0},
)
rows = result.fetchall()
-
+
assert len(rows) == 0
-
+
await storage.close()
@pytest.mark.asyncio
async def test_operations_after_close_raise_error(self, temp_db_path: Path) -> None:
"""Test that operations after close work properly."""
default_config = EventLogConfig()
- storage = AsyncSQLiteEventStorage(db_path=temp_db_path, batch_size=default_config.batch_size, batch_timeout_ms=default_config.batch_timeout_ms, debounce_ms=default_config.debounce_ms, max_age_ms=default_config.max_age_ms)
+ storage = AsyncSQLiteEventStorage(
+ db_path=temp_db_path,
+ batch_size=default_config.batch_size,
+ batch_timeout_ms=default_config.batch_timeout_ms,
+ debounce_ms=default_config.debounce_ms,
+ max_age_ms=default_config.max_age_ms,
+ )
await storage.start()
await storage.close()
-
+
# These should not raise errors since we're not using the public API
assert storage._closed is True
assert storage._engine is not None # Engine should still exist but be disposed
@@ -325,18 +422,32 @@ class TestAsyncSQLiteEventStorage:
async def test_multiple_close_calls_safe(self, temp_db_path: Path) -> None:
"""Test that multiple close calls are safe."""
default_config = EventLogConfig()
- storage = AsyncSQLiteEventStorage(db_path=temp_db_path, batch_size=default_config.batch_size, batch_timeout_ms=default_config.batch_timeout_ms, debounce_ms=default_config.debounce_ms, max_age_ms=default_config.max_age_ms)
+ storage = AsyncSQLiteEventStorage(
+ db_path=temp_db_path,
+ batch_size=default_config.batch_size,
+ batch_timeout_ms=default_config.batch_timeout_ms,
+ debounce_ms=default_config.debounce_ms,
+ max_age_ms=default_config.max_age_ms,
+ )
await storage.start()
await storage.close()
await storage.close() # Should not raise an error
@pytest.mark.asyncio
- async def test_json_data_types(self, temp_db_path: Path, sample_node_id: NodeId) -> None:
+ async def test_json_data_types(
+ self, temp_db_path: Path, sample_node_id: NodeId
+ ) -> None:
"""Test that various JSON data types are handled correctly."""
default_config = EventLogConfig()
- storage = AsyncSQLiteEventStorage(db_path=temp_db_path, batch_size=default_config.batch_size, batch_timeout_ms=default_config.batch_timeout_ms, debounce_ms=default_config.debounce_ms, max_age_ms=default_config.max_age_ms)
+ storage = AsyncSQLiteEventStorage(
+ db_path=temp_db_path,
+ batch_size=default_config.batch_size,
+ batch_timeout_ms=default_config.batch_timeout_ms,
+ debounce_ms=default_config.debounce_ms,
+ max_age_ms=default_config.max_age_ms,
+ )
await storage.start()
-
+
# Test various JSON data types
test_data = {
"string": "test string",
@@ -346,71 +457,81 @@ class TestAsyncSQLiteEventStorage:
"null": None,
"array": [1, 2, 3, "four"],
"object": {"nested": "value", "deep": {"deeper": "nested"}},
- "unicode": "测试 🚀"
+ "unicode": "测试 🚀",
}
-
+
assert storage._engine is not None
async with AsyncSession(storage._engine) as session:
await session.execute(
- text("INSERT INTO events (origin, event_type, event_id, event_data) VALUES (:origin, :event_type, :event_id, :event_data)"),
+ text(
+ "INSERT INTO events (origin, event_type, event_id, event_data) VALUES (:origin, :event_type, :event_id, :event_data)"
+ ),
{
"origin": sample_node_id,
"event_type": "complex_event",
"event_id": str(uuid4()),
- "event_data": json.dumps(test_data)
- }
+ "event_data": json.dumps(test_data),
+ },
)
await session.commit()
-
+
# Query back and verify data integrity
assert storage._engine is not None
async with AsyncSession(storage._engine) as session:
result = await session.execute(
text("SELECT event_data FROM events WHERE event_type = :event_type"),
- {"event_type": "complex_event"}
+ {"event_type": "complex_event"},
)
rows = result.fetchall()
-
+
assert len(rows) == 1
raw_json = cast(str, rows[0][0])
retrieved_data = _load_json_data(raw_json)
assert retrieved_data == test_data
-
+
await storage.close()
@pytest.mark.asyncio
async def test_concurrent_inserts(self, temp_db_path: Path) -> None:
"""Test concurrent inserts maintain rowid ordering."""
default_config = EventLogConfig()
- storage = AsyncSQLiteEventStorage(db_path=temp_db_path, batch_size=default_config.batch_size, batch_timeout_ms=default_config.batch_timeout_ms, debounce_ms=default_config.debounce_ms, max_age_ms=default_config.max_age_ms)
+ storage = AsyncSQLiteEventStorage(
+ db_path=temp_db_path,
+ batch_size=default_config.batch_size,
+ batch_timeout_ms=default_config.batch_timeout_ms,
+ debounce_ms=default_config.debounce_ms,
+ max_age_ms=default_config.max_age_ms,
+ )
await storage.start()
-
+
async def insert_batch(origin_id: str, batch_id: int, count: int) -> None:
assert storage._engine is not None
async with AsyncSession(storage._engine) as session:
for i in range(count):
await session.execute(
- text("INSERT INTO events (origin, event_type, event_id, event_data) VALUES (:origin, :event_type, :event_id, :event_data)"),
+ text(
+ "INSERT INTO events (origin, event_type, event_id, event_data) VALUES (:origin, :event_type, :event_id, :event_data)"
+ ),
{
"origin": origin_id,
"event_type": f"batch_{batch_id}_event_{i}",
"event_id": str(uuid4()),
- "event_data": json.dumps({"batch": batch_id, "item": i})
- }
+ "event_data": json.dumps({"batch": batch_id, "item": i}),
+ },
)
await session.commit()
-
+
# Run multiple concurrent insert batches
origin1 = str(uuid4())
origin2 = str(uuid4())
origin3 = str(uuid4())
-
+
await asyncio.gather(
insert_batch(origin1, 1, 5),
insert_batch(origin2, 2, 5),
- insert_batch(origin3, 3, 5)
+ insert_batch(origin3, 3, 5),
)
-
+
# Verify all records were inserted and rowid is sequential
assert storage._engine is not None
async with AsyncSession(storage._engine) as session:
@@ -418,22 +539,30 @@ class TestAsyncSQLiteEventStorage:
text("SELECT rowid, origin, event_data FROM events ORDER BY rowid")
)
rows = result.fetchall()
-
+
assert len(rows) == 15 # 3 batches * 5 records each
-
+
# Verify rowid sequence is maintained
for i, row in enumerate(rows):
assert row[0] == i + 1 # rowid should be sequential
-
+
await storage.close()
@pytest.mark.asyncio
- async def test_chunk_generated_event_serialization(self, temp_db_path: Path, sample_node_id: NodeId) -> None:
+ async def test_chunk_generated_event_serialization(
+ self, temp_db_path: Path, sample_node_id: NodeId
+ ) -> None:
"""Test that ChunkGenerated event with nested types can be serialized and deserialized correctly."""
default_config = EventLogConfig()
- storage = AsyncSQLiteEventStorage(db_path=temp_db_path, batch_size=default_config.batch_size, batch_timeout_ms=default_config.batch_timeout_ms, debounce_ms=default_config.debounce_ms, max_age_ms=default_config.max_age_ms)
+ storage = AsyncSQLiteEventStorage(
+ db_path=temp_db_path,
+ batch_size=default_config.batch_size,
+ batch_timeout_ms=default_config.batch_timeout_ms,
+ debounce_ms=default_config.debounce_ms,
+ max_age_ms=default_config.max_age_ms,
+ )
await storage.start()
-
+
# Create a ChunkGenerated event with nested TokenChunk
command_id = CommandId()
token_chunk = TokenChunk(
@@ -443,33 +572,30 @@ class TestAsyncSQLiteEventStorage:
chunk_type=ChunkType.token,
command_id=command_id,
idx=0,
- model="test-model"
- )
-
- chunk_generated_event = ChunkGenerated(
- command_id=command_id,
- chunk=token_chunk
+ model="test-model",
)
-
+
+ chunk_generated_event = ChunkGenerated(command_id=command_id, chunk=token_chunk)
+
# Store the event using the storage API
await storage.append_events([chunk_generated_event], sample_node_id)
-
+
# Wait for batch to be written
await asyncio.sleep(0.5)
-
+
# Retrieve the event
events = await storage.get_events_since(0)
-
+
# Verify we got the event back
assert len(events) == 1
retrieved_event_wrapper = events[0]
assert retrieved_event_wrapper.origin == sample_node_id
-
+
# Verify the event was deserialized correctly
retrieved_event = retrieved_event_wrapper.event
assert isinstance(retrieved_event, ChunkGenerated)
assert retrieved_event.command_id == command_id
-
+
# Verify the nested chunk was deserialized correctly
retrieved_chunk = retrieved_event.chunk
assert isinstance(retrieved_chunk, TokenChunk)
@@ -477,10 +603,10 @@ class TestAsyncSQLiteEventStorage:
assert retrieved_chunk.command_id == command_id
assert retrieved_chunk.idx == 0
assert retrieved_chunk.model == "test-model"
-
+
# Verify the chunk data
assert retrieved_chunk.text == "Hello, world!"
assert retrieved_chunk.token_id == 42
assert retrieved_chunk.finish_reason == "stop"
-
- await storage.close()
\ No newline at end of file
+
+ await storage.close()
diff --git a/src/exo/shared/topology.py b/src/exo/shared/topology.py
index 9322c721..a3825a27 100644
--- a/src/exo/shared/topology.py
+++ b/src/exo/shared/topology.py
@@ -53,7 +53,7 @@ class Topology(TopologyProto):
rx_id = self._graph.add_node(node)
self._node_id_to_rx_id_map[node.node_id] = rx_id
self._rx_id_to_node_id_map[rx_id] = node.node_id
-
+
def set_master_node_id(self, node_id: NodeId) -> None:
self.master_node_id = node_id
@@ -64,8 +64,8 @@ class Topology(TopologyProto):
return connection in self._edge_id_to_rx_id_map
def add_connection(
- self,
- connection: Connection,
+ self,
+ connection: Connection,
) -> None:
if connection.local_node_id not in self._node_id_to_rx_id_map:
self.add_node(Node(node_id=connection.local_node_id))
@@ -82,7 +82,7 @@ class Topology(TopologyProto):
yield from (self._graph[i] for i in self._graph.node_indices())
def list_connections(self) -> Iterable[Connection]:
- for (_, _, connection) in self._graph.weighted_edge_list():
+ for _, _, connection in self._graph.weighted_edge_list():
yield connection
def get_node_profile(self, node_id: NodeId) -> NodePerformanceProfile | None:
@@ -91,7 +91,7 @@ class Topology(TopologyProto):
return self._graph.get_node_data(rx_idx).node_profile
except KeyError:
return None
-
+
def get_node_multiaddr(self, node_id: NodeId) -> Multiaddr:
for connection in self.list_connections():
if connection.local_node_id == node_id:
@@ -99,8 +99,10 @@ class Topology(TopologyProto):
if connection.send_back_node_id == node_id:
return connection.send_back_multiaddr
raise ValueError(f"Node {node_id} is not connected to any other nodes")
-
- def update_node_profile(self, node_id: NodeId, node_profile: NodePerformanceProfile) -> None:
+
+ def update_node_profile(
+ self, node_id: NodeId, node_profile: NodePerformanceProfile
+ ) -> None:
rx_idx = self._node_id_to_rx_id_map[node_id]
self._graph[rx_idx].node_profile = node_profile
@@ -108,7 +110,9 @@ class Topology(TopologyProto):
rx_idx = self._edge_id_to_rx_id_map[connection]
self._graph.update_edge_by_index(rx_idx, connection)
- def get_connection_profile(self, connection: Connection) -> ConnectionProfile | None:
+ def get_connection_profile(
+ self, connection: Connection
+ ) -> ConnectionProfile | None:
try:
rx_idx = self._edge_id_to_rx_id_map[connection]
return self._graph.get_edge_data_by_index(rx_idx).connection_profile
@@ -128,14 +132,18 @@ class Topology(TopologyProto):
# Determine the reference node from which reachability is calculated.
# Prefer a master node if the topology knows one; otherwise fall back to
# the local end of the connection being removed.
- reference_node_id: NodeId = self.master_node_id if self.master_node_id is not None else connection.local_node_id
+ reference_node_id: NodeId = (
+ self.master_node_id
+ if self.master_node_id is not None
+ else connection.local_node_id
+ )
orphan_node_ids = self._get_orphan_node_ids(reference_node_id, connection)
for orphan_node_id in orphan_node_ids:
orphan_node_rx_id = self._node_id_to_rx_id_map[orphan_node_id]
self._graph.remove_node(orphan_node_rx_id)
del self._node_id_to_rx_id_map[orphan_node_id]
del self._rx_id_to_node_id_map[orphan_node_rx_id]
-
+
self._graph.remove_edge_from_index(rx_idx)
del self._edge_id_to_rx_id_map[connection]
if rx_idx in self._rx_id_to_node_id_map:
@@ -149,7 +157,7 @@ class Topology(TopologyProto):
cycles.append(cycle)
return cycles
-
+
def get_subgraph_from_nodes(self, nodes: list[Node]) -> "Topology":
node_idxs = [node.node_id for node in nodes]
rx_idxs = [self._node_id_to_rx_id_map[idx] for idx in node_idxs]
@@ -157,7 +165,10 @@ class Topology(TopologyProto):
for rx_idx in rx_idxs:
topology.add_node(self._graph[rx_idx])
for connection in self.list_connections():
- if connection.local_node_id in node_idxs and connection.send_back_node_id in node_idxs:
+ if (
+ connection.local_node_id in node_idxs
+ and connection.send_back_node_id in node_idxs
+ ):
topology.add_connection(connection)
return topology
@@ -176,16 +187,18 @@ class Topology(TopologyProto):
if not has_tb:
return False
return True
-
+
def _is_bridge(self, connection: Connection) -> bool:
"""Check if removing this connection will orphan any nodes from the master."""
if self.master_node_id is None:
return False
-
+
orphan_node_ids = self._get_orphan_node_ids(self.master_node_id, connection)
return len(orphan_node_ids) > 0
- def _get_orphan_node_ids(self, master_node_id: NodeId, connection: Connection) -> list[NodeId]:
+ def _get_orphan_node_ids(
+ self, master_node_id: NodeId, connection: Connection
+ ) -> list[NodeId]:
"""Return node_ids that become unreachable from `master_node_id` once `connection` is removed.
A node is considered *orphaned* if there exists **no directed path** from
@@ -215,4 +228,8 @@ class Topology(TopologyProto):
# Every existing node index not reachable is orphaned.
orphan_rx_ids = set(graph_copy.node_indices()) - reachable_rx_ids
- return [self._rx_id_to_node_id_map[rx_id] for rx_id in orphan_rx_ids if rx_id in self._rx_id_to_node_id_map]
+ return [
+ self._rx_id_to_node_id_map[rx_id]
+ for rx_id in orphan_rx_ids
+ if rx_id in self._rx_id_to_node_id_map
+ ]
diff --git a/src/exo/shared/types/api.py b/src/exo/shared/types/api.py
index fc05d160..22870e63 100644
--- a/src/exo/shared/types/api.py
+++ b/src/exo/shared/types/api.py
@@ -21,10 +21,12 @@ class ModelListModel(BaseModel):
context_length: int = Field(default=0)
tags: List[str] = Field(default=[])
+
class ModelList(BaseModel):
object: str = "list"
data: List[ModelListModel]
+
class ChatCompletionMessage(BaseModel):
role: Literal["system", "user", "assistant", "developer", "tool", "function"]
content: str | None = None
@@ -86,7 +88,6 @@ class Usage(BaseModel):
completion_tokens_details: CompletionTokensDetails | None = None
-
class ChatCompletionResponse(BaseModel):
id: str
object: Literal["chat.completion"] = "chat.completion"
@@ -118,19 +119,23 @@ class ChatCompletionTaskParams(BaseModel):
parallel_tool_calls: bool | None = None
user: str | None = None
+
class CreateInstanceTaskParams(BaseModel):
# TODO: in future the user could specify a specific Instance, not just a model_id
model_id: str
+
class DeleteInstanceTaskParams(BaseModel):
instance_id: str
+
class CreateInstanceResponse(BaseModel):
message: str
command_id: CommandId
model_meta: ModelMetadata
instance_id: InstanceId
+
class DeleteInstanceResponse(BaseModel):
message: str
command_id: CommandId
diff --git a/src/exo/shared/types/common.py b/src/exo/shared/types/common.py
index c949712b..bc7cd127 100644
--- a/src/exo/shared/types/common.py
+++ b/src/exo/shared/types/common.py
@@ -12,9 +12,7 @@ class ID(str):
@classmethod
def __get_pydantic_core_schema__(
- cls,
- _source: type[Any],
- handler: GetCoreSchemaHandler
+ cls, _source: type[Any], handler: GetCoreSchemaHandler
) -> core_schema.CoreSchema:
# Re‑use the already‑defined schema for `str`
return handler.generate_schema(str)
@@ -41,4 +39,3 @@ class Host(BaseModel):
if not (0 <= v <= 65535):
raise ValueError("Port must be between 0 and 65535")
return v
-
diff --git a/src/exo/shared/types/events/_events.py b/src/exo/shared/types/events/_events.py
index b61be0e5..c59a2df1 100644
--- a/src/exo/shared/types/events/_events.py
+++ b/src/exo/shared/types/events/_events.py
@@ -40,6 +40,7 @@ class EventId(ID):
# Event base-class boilerplate (you should basically never touch these)
# Only very specialised registry or serialisation/deserialization logic might need know about these
+
class _EventType(str, Enum):
"""
Here are all the unique kinds of events that can be sent over the network.
@@ -102,8 +103,10 @@ class _BaseEvent[T: _EventType](BaseModel):
"""
return True
+
_E = TypeVar("_E", bound=_BaseEvent[Any])
+
def no_op_event(cls: type[_E]) -> type[_E]:
"""Decorator to mark an event class as a *no-op*.
@@ -115,11 +118,14 @@ def no_op_event(cls: type[_E]) -> type[_E]:
cls.__no_apply__ = True # Used by the apply layer to identify no-op events
return cls
+
+
@no_op_event
class Heartbeat(_BaseEvent[_EventType.Heartbeat]):
event_type: Literal[_EventType.Heartbeat] = _EventType.Heartbeat
node_id: NodeId
+
class TaskCreated(_BaseEvent[_EventType.TaskCreated]):
event_type: Literal[_EventType.TaskCreated] = _EventType.TaskCreated
task_id: TaskId
@@ -165,12 +171,16 @@ class InstanceDeleted(_BaseEvent[_EventType.InstanceDeleted]):
class InstanceReplacedAtomically(_BaseEvent[_EventType.InstanceReplacedAtomically]):
- event_type: Literal[_EventType.InstanceReplacedAtomically] = _EventType.InstanceReplacedAtomically
+ event_type: Literal[_EventType.InstanceReplacedAtomically] = (
+ _EventType.InstanceReplacedAtomically
+ )
instance_to_replace: InstanceId
new_instance_id: InstanceId
+
# TODO: RunnerCreated
+
class RunnerStatusUpdated(_BaseEvent[_EventType.RunnerStatusUpdated]):
event_type: Literal[_EventType.RunnerStatusUpdated] = _EventType.RunnerStatusUpdated
runner_id: RunnerId
@@ -183,7 +193,9 @@ class RunnerDeleted(_BaseEvent[_EventType.RunnerDeleted]):
class NodePerformanceMeasured(_BaseEvent[_EventType.NodePerformanceMeasured]):
- event_type: Literal[_EventType.NodePerformanceMeasured] = _EventType.NodePerformanceMeasured
+ event_type: Literal[_EventType.NodePerformanceMeasured] = (
+ _EventType.NodePerformanceMeasured
+ )
node_id: NodeId
node_profile: NodePerformanceProfile
@@ -200,22 +212,28 @@ class ChunkGenerated(_BaseEvent[_EventType.ChunkGenerated]):
command_id: CommandId
chunk: GenerationChunk
+
class TopologyNodeCreated(_BaseEvent[_EventType.TopologyNodeCreated]):
event_type: Literal[_EventType.TopologyNodeCreated] = _EventType.TopologyNodeCreated
node_id: NodeId
role: Literal["MASTER", "REPLICA"]
+
class TopologyEdgeCreated(_BaseEvent[_EventType.TopologyEdgeCreated]):
event_type: Literal[_EventType.TopologyEdgeCreated] = _EventType.TopologyEdgeCreated
edge: Connection
-class TopologyEdgeReplacedAtomically(_BaseEvent[_EventType.TopologyEdgeReplacedAtomically]):
+class TopologyEdgeReplacedAtomically(
+ _BaseEvent[_EventType.TopologyEdgeReplacedAtomically]
+):
"""
TODO: delete this????
"""
- event_type: Literal[_EventType.TopologyEdgeReplacedAtomically] = _EventType.TopologyEdgeReplacedAtomically
+ event_type: Literal[_EventType.TopologyEdgeReplacedAtomically] = (
+ _EventType.TopologyEdgeReplacedAtomically
+ )
edge: Connection
edge_profile: ConnectionProfile
@@ -262,30 +280,34 @@ def _check_event_type_consistency():
for cls in union_classes: # pyright: ignore[reportAny]
assert issubclass(cls, object), (
f"{get_error_reporting_message()}",
- f"The class {cls} is NOT a subclass of {object}."
+ f"The class {cls} is NOT a subclass of {object}.",
)
# ensure the first base parameter is ALWAYS _BaseEvent
base_cls = list(types.get_original_bases(cls))
- assert len(base_cls) >= 1 and issubclass(base_cls[0], object) \
- and issubclass(base_cls[0], _BaseEvent), (
+ assert (
+ len(base_cls) >= 1
+ and issubclass(base_cls[0], object)
+ and issubclass(base_cls[0], _BaseEvent)
+ ), (
f"{get_error_reporting_message()}",
- f"The class {cls} does NOT inherit from {_BaseEvent} {get_origin(base_cls[0])}."
+ f"The class {cls} does NOT inherit from {_BaseEvent} {get_origin(base_cls[0])}.",
)
# grab type hints and extract the right values from it
cls_hints = get_type_hints(cls)
- assert "event_type" in cls_hints and \
- get_origin(cls_hints["event_type"]) is Literal, ( # pyright: ignore[reportAny]
+ assert (
+ "event_type" in cls_hints and get_origin(cls_hints["event_type"]) is Literal
+ ), ( # pyright: ignore[reportAny]
f"{get_error_reporting_message()}",
- f"The class {cls} is missing a {Literal}-annotated `event_type` field."
+ f"The class {cls} is missing a {Literal}-annotated `event_type` field.",
)
# make sure the value is an instance of `_EventType`
enum_value = list(get_args(cls_hints["event_type"]))
assert len(enum_value) == 1 and isinstance(enum_value[0], _EventType), (
f"{get_error_reporting_message()}",
- f"The `event_type` of {cls} has a non-{_EventType} literal-type."
+ f"The `event_type` of {cls} has a non-{_EventType} literal-type.",
)
union_enum_values.append(enum_value[0])
@@ -293,12 +315,12 @@ def _check_event_type_consistency():
for m in member_enum_values:
assert m in union_enum_values, (
f"{get_error_reporting_message()}",
- f"There is no event-type registered for {m} in {_Event}."
+ f"There is no event-type registered for {m} in {_Event}.",
)
union_enum_values.remove(m)
assert len(union_enum_values) == 0, (
f"{get_error_reporting_message()}",
- f"The following events have multiple event types defined in {_Event}: {union_enum_values}."
+ f"The following events have multiple event types defined in {_Event}: {union_enum_values}.",
)
diff --git a/src/exo/shared/types/events/chunks.py b/src/exo/shared/types/events/chunks.py
index ebf68ace..7a69ae5c 100644
--- a/src/exo/shared/types/events/chunks.py
+++ b/src/exo/shared/types/events/chunks.py
@@ -31,6 +31,7 @@ class ImageChunk(BaseChunk[ChunkType.image]):
chunk_type: Literal[ChunkType.image] = Field(default=ChunkType.image, frozen=True)
data: bytes
+
GenerationChunk = Annotated[TokenChunk | ImageChunk, Field(discriminator="chunk_type")]
GenerationChunkTypeAdapter: TypeAdapter[GenerationChunk] = TypeAdapter(GenerationChunk)
@@ -41,8 +42,8 @@ GenerationChunkTypeAdapter: TypeAdapter[GenerationChunk] = TypeAdapter(Generatio
# my_chunk: dict[str, Any] = TokenChunk(
# task_id=TaskId('nicerid'),
# idx=0,
- # text='hello',
- # token_id=12,
+# text='hello',
+# token_id=12,
# chunk_type=ChunkType.token,
# model='llama-3.1',
# ).model_dump()
diff --git a/src/exo/shared/types/events/commands.py b/src/exo/shared/types/events/commands.py
index 8f60f18b..7469e1fa 100644
--- a/src/exo/shared/types/events/commands.py
+++ b/src/exo/shared/types/events/commands.py
@@ -45,8 +45,11 @@ class TaskFinishedCommand(_BaseCommand[CommandType.TASK_FINISHED]):
Command = Annotated[
- ChatCompletionCommand | CreateInstanceCommand | DeleteInstanceCommand | TaskFinishedCommand,
- Field(discriminator="command_type")
+ ChatCompletionCommand
+ | CreateInstanceCommand
+ | DeleteInstanceCommand
+ | TaskFinishedCommand,
+ Field(discriminator="command_type"),
]
CommandParser: TypeAdapter[Command] = TypeAdapter(Command)
diff --git a/src/exo/shared/types/events/components.py b/src/exo/shared/types/events/components.py
index b9ef7620..d0764b85 100644
--- a/src/exo/shared/types/events/components.py
+++ b/src/exo/shared/types/events/components.py
@@ -1,4 +1,4 @@
-# components.py defines the small event functions, adapters etc.
+# components.py defines the small event functions, adapters etc.
# this name could probably be improved.
from typing import (
@@ -32,6 +32,5 @@ class EventFromEventLog[T: Event](BaseModel):
raise ValueError("Invalid Event: Origin ID Does Not Match")
-
type Apply = Callable[[State, Event], State]
type ApplyFromEventLog = Callable[[State, EventFromEventLog[Event]], State]
diff --git a/src/exo/shared/types/graphs/pydantic.py b/src/exo/shared/types/graphs/pydantic.py
index 2ff9e557..ce2afabb 100644
--- a/src/exo/shared/types/graphs/pydantic.py
+++ b/src/exo/shared/types/graphs/pydantic.py
@@ -5,4 +5,4 @@ from pydantic import BaseModel
class PydanticGraph(BaseModel):
vertices: List[Any]
- edges: List[Any]
\ No newline at end of file
+ edges: List[Any]
diff --git a/src/exo/shared/types/multiaddr.py b/src/exo/shared/types/multiaddr.py
index 7cbdadec..23cf55ae 100644
--- a/src/exo/shared/types/multiaddr.py
+++ b/src/exo/shared/types/multiaddr.py
@@ -7,13 +7,13 @@ from pydantic import BaseModel, computed_field, field_serializer, field_validato
class Multiaddr(BaseModel):
address: str
-
+
PATTERNS: ClassVar[list[str]] = [
- r'^/ip4/(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})(/tcp/(\d{1,5}))?(/p2p/[A-Za-z0-9]+)?$',
- r'^/ip6/([0-9a-fA-F:]+)(/tcp/(\d{1,5}))?(/p2p/[A-Za-z0-9]+)?$',
- r'^/dns[46]?/([a-zA-Z0-9.-]+)(/tcp/(\d{1,5}))?(/p2p/[A-Za-z0-9]+)?$',
+ r"^/ip4/(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})(/tcp/(\d{1,5}))?(/p2p/[A-Za-z0-9]+)?$",
+ r"^/ip6/([0-9a-fA-F:]+)(/tcp/(\d{1,5}))?(/p2p/[A-Za-z0-9]+)?$",
+ r"^/dns[46]?/([a-zA-Z0-9.-]+)(/tcp/(\d{1,5}))?(/p2p/[A-Za-z0-9]+)?$",
]
-
+
@field_validator("address")
@classmethod
def validate_format(cls, v: str) -> str:
@@ -23,34 +23,38 @@ class Multiaddr(BaseModel):
"Expected format like /ip4/127.0.0.1/tcp/4001 or /dns/example.com/tcp/443"
)
return v
-
+
@computed_field
@property
def address_type(self) -> str:
for pattern in self.PATTERNS:
if re.match(pattern, self.address):
- return pattern.split('/')[1]
+ return pattern.split("/")[1]
raise ValueError(f"Invalid multiaddr format: {self.address}")
-
+
@property
def ipv6_address(self) -> IPv6Address:
- match = re.match(r'^/ip6/([0-9a-fA-F:]+)', self.address)
+ match = re.match(r"^/ip6/([0-9a-fA-F:]+)", self.address)
if not match:
- raise ValueError(f"Invalid multiaddr format: {self.address}. Expected format like /ip6/::1/tcp/4001")
+ raise ValueError(
+ f"Invalid multiaddr format: {self.address}. Expected format like /ip6/::1/tcp/4001"
+ )
return IPv6Address(match.group(1))
-
+
@property
def ipv4_address(self) -> IPv4Address:
- match = re.match(r'^/ip4/(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})', self.address)
+ match = re.match(r"^/ip4/(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})", self.address)
if not match:
- raise ValueError(f"Invalid multiaddr format: {self.address}. Expected format like /ip4/127.0.0.1/tcp/4001")
+ raise ValueError(
+ f"Invalid multiaddr format: {self.address}. Expected format like /ip4/127.0.0.1/tcp/4001"
+ )
return IPv4Address(match.group(1))
@computed_field
@property
def ip_address(self) -> IPv4Address | IPv6Address:
- return self.ipv4_address if self.address_type == 'ip4' else self.ipv6_address
-
+ return self.ipv4_address if self.address_type == "ip4" else self.ipv6_address
+
@field_serializer("ip_address")
def serialize_ipv4_address(self, value: IPv4Address | IPv6Address) -> str:
return str(value)
@@ -58,11 +62,12 @@ class Multiaddr(BaseModel):
@computed_field
@property
def port(self) -> int:
- match = re.search(r'/tcp/(\d{1,5})', self.address)
+ match = re.search(r"/tcp/(\d{1,5})", self.address)
if not match:
- raise ValueError(f"Invalid multiaddr format: {self.address}. Expected format like /ip4/127.0.0.1/tcp/4001")
+ raise ValueError(
+ f"Invalid multiaddr format: {self.address}. Expected format like /ip4/127.0.0.1/tcp/4001"
+ )
return int(match.group(1))
-
def __str__(self) -> str:
return self.address
diff --git a/src/exo/shared/types/request.py b/src/exo/shared/types/request.py
index 0e8d6b4c..d471be8b 100644
--- a/src/exo/shared/types/request.py
+++ b/src/exo/shared/types/request.py
@@ -12,12 +12,15 @@ class ChatCompletionCommand(BaseModel):
command_id: CommandId
command_params: ChatCompletionTaskParams
+
class CreateInstanceCommand(BaseModel):
command_id: CommandId
command_params: CreateInstanceTaskParams
+
class DeleteInstanceCommand(BaseModel):
command_id: CommandId
command_params: DeleteInstanceTaskParams
+
type Command = ChatCompletionCommand | CreateInstanceCommand | DeleteInstanceCommand
diff --git a/src/exo/shared/types/state.py b/src/exo/shared/types/state.py
index bf9eca8f..368400df 100644
--- a/src/exo/shared/types/state.py
+++ b/src/exo/shared/types/state.py
@@ -17,6 +17,7 @@ def _encode_topology(topo: "Topology") -> dict[str, Any]: # noqa: D401
return topo.to_snapshot().model_dump()
+
class State(BaseModel):
"""Global system state.
diff --git a/src/exo/shared/types/tasks.py b/src/exo/shared/types/tasks.py
index ea609f28..58f4b67f 100644
--- a/src/exo/shared/types/tasks.py
+++ b/src/exo/shared/types/tasks.py
@@ -34,4 +34,5 @@ class ChatCompletionTask(BaseModel):
error_type: Optional[str] = Field(default=None)
error_message: Optional[str] = Field(default=None)
+
Task = Annotated[ChatCompletionTask, Field(discriminator="task_type")]
diff --git a/src/exo/shared/types/topology.py b/src/exo/shared/types/topology.py
index fc87b484..98f1d29c 100644
--- a/src/exo/shared/types/topology.py
+++ b/src/exo/shared/types/topology.py
@@ -31,14 +31,17 @@ class Connection(BaseModel):
if not isinstance(other, Connection):
raise ValueError("Cannot compare Connection with non-Connection")
return (
- self.local_node_id == other.local_node_id
- and self.send_back_node_id == other.send_back_node_id
- and self.local_multiaddr.ip_address == other.local_multiaddr.ip_address
- and self.send_back_multiaddr.ip_address == other.send_back_multiaddr.ip_address
+ self.local_node_id == other.local_node_id
+ and self.send_back_node_id == other.send_back_node_id
+ and self.local_multiaddr.ip_address == other.local_multiaddr.ip_address
+ and self.send_back_multiaddr.ip_address
+ == other.send_back_multiaddr.ip_address
)
-
+
def is_thunderbolt(self) -> bool:
- return str(self.local_multiaddr.ip_address).startswith('169.254') and str(self.send_back_multiaddr.ip_address).startswith('169.254')
+ return str(self.local_multiaddr.ip_address).startswith("169.254") and str(
+ self.send_back_multiaddr.ip_address
+ ).startswith("169.254")
class Node(BaseModel):
@@ -50,15 +53,17 @@ class TopologyProto(Protocol):
def add_node(self, node: Node) -> None: ...
def add_connection(
- self,
- connection: Connection,
+ self,
+ connection: Connection,
) -> None: ...
def list_nodes(self) -> Iterable[Node]: ...
def list_connections(self) -> Iterable[Connection]: ...
- def update_node_profile(self, node_id: NodeId, node_profile: NodePerformanceProfile) -> None: ...
+ def update_node_profile(
+ self, node_id: NodeId, node_profile: NodePerformanceProfile
+ ) -> None: ...
def update_connection_profile(self, connection: Connection) -> None: ...
@@ -68,6 +73,8 @@ class TopologyProto(Protocol):
def get_node_profile(self, node_id: NodeId) -> NodePerformanceProfile | None: ...
- def get_connection_profile(self, connection: Connection) -> ConnectionProfile | None: ...
+ def get_connection_profile(
+ self, connection: Connection
+ ) -> ConnectionProfile | None: ...
def get_cycles(self) -> list[list[Node]]: ...
diff --git a/src/exo/shared/types/worker/commands_runner.py b/src/exo/shared/types/worker/commands_runner.py
index 19f96f68..be3b27c5 100644
--- a/src/exo/shared/types/worker/commands_runner.py
+++ b/src/exo/shared/types/worker/commands_runner.py
@@ -103,8 +103,13 @@ class ErrorResponse(BaseRunnerResponse[RunnerResponseType.ErrorResponse]):
error_message: str
traceback: str
+
RunnerResponse = Annotated[
- InitializedResponse | GenerationResponse | PrintResponse | FinishedResponse | ErrorResponse,
+ InitializedResponse
+ | GenerationResponse
+ | PrintResponse
+ | FinishedResponse
+ | ErrorResponse,
Field(discriminator="type"),
]
RunnerResponseTypeAdapter: TypeAdapter[RunnerResponse] = TypeAdapter(RunnerResponse)
diff --git a/src/exo/shared/types/worker/common.py b/src/exo/shared/types/worker/common.py
index 2f72de6f..37502167 100644
--- a/src/exo/shared/types/worker/common.py
+++ b/src/exo/shared/types/worker/common.py
@@ -15,11 +15,12 @@ class NodeStatus(str, Enum):
Idle = "Idle"
Running = "Running"
+
class RunnerError(Exception):
- """Exception raised when the runner process encounters an error."""
-
- def __init__(self, error_type: str, error_message: str, traceback: str):
- self.error_type = error_type
- self.error_message = error_message
- self.traceback = traceback
- super().__init__(f"{error_type}: {error_message}. Traceback: {traceback}")
\ No newline at end of file
+ """Exception raised when the runner process encounters an error."""
+
+ def __init__(self, error_type: str, error_message: str, traceback: str):
+ self.error_type = error_type
+ self.error_message = error_message
+ self.traceback = traceback
+ super().__init__(f"{error_type}: {error_message}. Traceback: {traceback}")
diff --git a/src/exo/shared/types/worker/downloads.py b/src/exo/shared/types/worker/downloads.py
index 8415dc55..54672205 100644
--- a/src/exo/shared/types/worker/downloads.py
+++ b/src/exo/shared/types/worker/downloads.py
@@ -33,15 +33,21 @@ class BaseDownloadProgress[DownloadStatusT: DownloadStatus](BaseModel):
class DownloadPending(BaseDownloadProgress[DownloadStatus.Pending]):
- download_status: Literal[DownloadStatus.Pending] = Field(default=DownloadStatus.Pending)
+ download_status: Literal[DownloadStatus.Pending] = Field(
+ default=DownloadStatus.Pending
+ )
class DownloadCompleted(BaseDownloadProgress[DownloadStatus.Completed]):
- download_status: Literal[DownloadStatus.Completed] = Field(default=DownloadStatus.Completed)
+ download_status: Literal[DownloadStatus.Completed] = Field(
+ default=DownloadStatus.Completed
+ )
class DownloadFailed(BaseDownloadProgress[DownloadStatus.Failed]):
- download_status: Literal[DownloadStatus.Failed] = Field(default=DownloadStatus.Failed)
+ download_status: Literal[DownloadStatus.Failed] = Field(
+ default=DownloadStatus.Failed
+ )
error_message: str
diff --git a/src/exo/shared/types/worker/instances.py b/src/exo/shared/types/worker/instances.py
index 9b0521c4..d44a0e54 100644
--- a/src/exo/shared/types/worker/instances.py
+++ b/src/exo/shared/types/worker/instances.py
@@ -13,6 +13,7 @@ class InstanceStatus(str, Enum):
ACTIVE = "ACTIVE"
INACTIVE = "INACTIVE"
+
class Instance(BaseModel):
instance_id: InstanceId
instance_type: InstanceStatus
diff --git a/src/exo/shared/types/worker/ops.py b/src/exo/shared/types/worker/ops.py
index b06dc0e1..386e2f4b 100644
--- a/src/exo/shared/types/worker/ops.py
+++ b/src/exo/shared/types/worker/ops.py
@@ -18,36 +18,56 @@ class RunnerOpType(str, Enum):
RUNNER_FAILED = "runner_failed"
CHAT_COMPLETION = "chat_completion"
+
RunnerOpT = TypeVar("RunnerOpT", bound=RunnerOpType)
+
class BaseRunnerOp(BaseModel, Generic[RunnerOpT]):
op_type: RunnerOpT
+
class AssignRunnerOp(BaseRunnerOp[Literal[RunnerOpType.ASSIGN_RUNNER]]):
- op_type: Literal[RunnerOpType.ASSIGN_RUNNER] = Field(default=RunnerOpType.ASSIGN_RUNNER, frozen=True)
+ op_type: Literal[RunnerOpType.ASSIGN_RUNNER] = Field(
+ default=RunnerOpType.ASSIGN_RUNNER, frozen=True
+ )
instance_id: InstanceId
runner_id: RunnerId
shard_metadata: ShardMetadata
hosts: list[Host]
+
class UnassignRunnerOp(BaseRunnerOp[Literal[RunnerOpType.UNASSIGN_RUNNER]]):
- op_type: Literal[RunnerOpType.UNASSIGN_RUNNER] = Field(default=RunnerOpType.UNASSIGN_RUNNER, frozen=True)
+ op_type: Literal[RunnerOpType.UNASSIGN_RUNNER] = Field(
+ default=RunnerOpType.UNASSIGN_RUNNER, frozen=True
+ )
runner_id: RunnerId
+
class RunnerUpOp(BaseRunnerOp[Literal[RunnerOpType.RUNNER_UP]]):
- op_type: Literal[RunnerOpType.RUNNER_UP] = Field(default=RunnerOpType.RUNNER_UP, frozen=True)
+ op_type: Literal[RunnerOpType.RUNNER_UP] = Field(
+ default=RunnerOpType.RUNNER_UP, frozen=True
+ )
runner_id: RunnerId
+
class RunnerDownOp(BaseRunnerOp[Literal[RunnerOpType.RUNNER_DOWN]]):
- op_type: Literal[RunnerOpType.RUNNER_DOWN] = Field(default=RunnerOpType.RUNNER_DOWN, frozen=True)
+ op_type: Literal[RunnerOpType.RUNNER_DOWN] = Field(
+ default=RunnerOpType.RUNNER_DOWN, frozen=True
+ )
runner_id: RunnerId
+
class RunnerFailedOp(BaseRunnerOp[Literal[RunnerOpType.RUNNER_FAILED]]):
- op_type: Literal[RunnerOpType.RUNNER_FAILED] = Field(default=RunnerOpType.RUNNER_FAILED, frozen=True)
+ op_type: Literal[RunnerOpType.RUNNER_FAILED] = Field(
+ default=RunnerOpType.RUNNER_FAILED, frozen=True
+ )
runner_id: RunnerId
+
class ExecuteTaskOp(BaseRunnerOp[Literal[RunnerOpType.CHAT_COMPLETION]]):
- op_type: Literal[RunnerOpType.CHAT_COMPLETION] = Field(default=RunnerOpType.CHAT_COMPLETION, frozen=True)
+ op_type: Literal[RunnerOpType.CHAT_COMPLETION] = Field(
+ default=RunnerOpType.CHAT_COMPLETION, frozen=True
+ )
runner_id: RunnerId
task: Task
@@ -62,5 +82,5 @@ RunnerOp = Annotated[
RunnerFailedOp,
ExecuteTaskOp,
],
- Field(discriminator="op_type")
-]
\ No newline at end of file
+ Field(discriminator="op_type"),
+]
diff --git a/src/exo/shared/types/worker/resource_monitor.py b/src/exo/shared/types/worker/resource_monitor.py
index 0bcdcfa4..8a1f3349 100644
--- a/src/exo/shared/types/worker/resource_monitor.py
+++ b/src/exo/shared/types/worker/resource_monitor.py
@@ -24,12 +24,16 @@ class MemoryResourceCollector(ResourceCollector):
class ResourceMonitor:
data_collectors: List[ResourceCollector]
- effect_handlers: Set[Callable[[SystemPerformanceProfile | MemoryPerformanceProfile], None]]
-
- async def _collect(self) -> list[SystemPerformanceProfile | MemoryPerformanceProfile]:
- tasks: list[Coroutine[None, None, SystemPerformanceProfile | MemoryPerformanceProfile]] = [
- collector.collect() for collector in self.data_collectors
- ]
+ effect_handlers: Set[
+ Callable[[SystemPerformanceProfile | MemoryPerformanceProfile], None]
+ ]
+
+ async def _collect(
+ self,
+ ) -> list[SystemPerformanceProfile | MemoryPerformanceProfile]:
+ tasks: list[
+ Coroutine[None, None, SystemPerformanceProfile | MemoryPerformanceProfile]
+ ] = [collector.collect() for collector in self.data_collectors]
return await asyncio.gather(*tasks)
async def collect(self) -> None:
diff --git a/src/exo/shared/types/worker/runners.py b/src/exo/shared/types/worker/runners.py
index 2abbc838..3bc70b5f 100644
--- a/src/exo/shared/types/worker/runners.py
+++ b/src/exo/shared/types/worker/runners.py
@@ -28,23 +28,40 @@ class BaseRunnerStatus(BaseModel, Generic[RunnerStatusTypeT]):
class DownloadingRunnerStatus(BaseRunnerStatus[RunnerStatusType.Downloading]):
- runner_status: Literal[RunnerStatusType.Downloading] = Field(default=RunnerStatusType.Downloading)
+ runner_status: Literal[RunnerStatusType.Downloading] = Field(
+ default=RunnerStatusType.Downloading
+ )
download_progress: DownloadProgress
+
class InactiveRunnerStatus(BaseRunnerStatus[RunnerStatusType.Inactive]):
- runner_status: Literal[RunnerStatusType.Inactive] = Field(default=RunnerStatusType.Inactive)
+ runner_status: Literal[RunnerStatusType.Inactive] = Field(
+ default=RunnerStatusType.Inactive
+ )
+
class StartingRunnerStatus(BaseRunnerStatus[RunnerStatusType.Starting]):
- runner_status: Literal[RunnerStatusType.Starting] = Field(default=RunnerStatusType.Starting)
+ runner_status: Literal[RunnerStatusType.Starting] = Field(
+ default=RunnerStatusType.Starting
+ )
+
class LoadedRunnerStatus(BaseRunnerStatus[RunnerStatusType.Loaded]):
- runner_status: Literal[RunnerStatusType.Loaded] = Field(default=RunnerStatusType.Loaded)
+ runner_status: Literal[RunnerStatusType.Loaded] = Field(
+ default=RunnerStatusType.Loaded
+ )
+
class RunningRunnerStatus(BaseRunnerStatus[RunnerStatusType.Running]):
- runner_status: Literal[RunnerStatusType.Running] = Field(default=RunnerStatusType.Running)
+ runner_status: Literal[RunnerStatusType.Running] = Field(
+ default=RunnerStatusType.Running
+ )
+
class FailedRunnerStatus(BaseRunnerStatus[RunnerStatusType.Failed]):
- runner_status: Literal[RunnerStatusType.Failed] = Field(default=RunnerStatusType.Failed)
+ runner_status: Literal[RunnerStatusType.Failed] = Field(
+ default=RunnerStatusType.Failed
+ )
error_message: str | None = None
@@ -57,9 +74,7 @@ RunnerStatus = Annotated[
| FailedRunnerStatus,
Field,
]
-RunnerStatusParser: TypeAdapter[RunnerStatus] = TypeAdapter(
- RunnerStatus
-)
+RunnerStatusParser: TypeAdapter[RunnerStatus] = TypeAdapter(RunnerStatus)
class ShardAssignments(BaseModel):
diff --git a/src/exo/shared/types/worker/shards.py b/src/exo/shared/types/worker/shards.py
index 62266652..d0602877 100644
--- a/src/exo/shared/types/worker/shards.py
+++ b/src/exo/shared/types/worker/shards.py
@@ -11,7 +11,9 @@ class PartitionStrategy(str, Enum):
pipeline = "pipeline"
-PartitionStrategyT = TypeVar("PartitionStrategyT", bound=PartitionStrategy, covariant=True)
+PartitionStrategyT = TypeVar(
+ "PartitionStrategyT", bound=PartitionStrategy, covariant=True
+)
class BaseShardMetadata(BaseModel, Generic[PartitionStrategyT]):
@@ -24,7 +26,7 @@ class BaseShardMetadata(BaseModel, Generic[PartitionStrategyT]):
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
@@ -34,7 +36,7 @@ class BaseShardMetadata(BaseModel, Generic[PartitionStrategyT]):
class PipelineShardMetadata(BaseShardMetadata[Literal[PartitionStrategy.pipeline]]):
"""
Pipeline parallelism shard meta.
-
+
Layers are represented as a half-open interval [start_layer, end_layer),
where start_layer is inclusive and end_layer is exclusive.
"""
@@ -49,21 +51,21 @@ class PipelineShardMetadata(BaseShardMetadata[Literal[PartitionStrategy.pipeline
@property
def is_first_layer(self) -> bool:
return self.start_layer == 0
-
+
@property
def is_last_layer(self) -> bool:
return self.end_layer == self.n_layers
def __hash__(self) -> int:
- return hash((self.model_meta.model_id, self.start_layer, self.end_layer, self.n_layers))
+ return hash(
+ (self.model_meta.model_id, self.start_layer, self.end_layer, self.n_layers)
+ )
ShardMetadata = Annotated[
PipelineShardMetadata, Field(discriminator="partition_strategy")
]
-ShardMetadataParser: TypeAdapter[ShardMetadata] = TypeAdapter(
- ShardMetadata
-)
+ShardMetadataParser: TypeAdapter[ShardMetadata] = TypeAdapter(ShardMetadata)
class ShardPlacement(BaseModel, Generic[PartitionStrategyT]):
diff --git a/src/exo/shared/utils.py b/src/exo/shared/utils.py
index df45ec6f..a819e7fb 100644
--- a/src/exo/shared/utils.py
+++ b/src/exo/shared/utils.py
@@ -20,15 +20,15 @@ class PeerId:
A libp2p peer identifier derived from a cryptographic public key.
Compatible with py-libp2p's PeerID interface.
"""
-
+
def __init__(self, peer_id_bytes: bytes) -> None:
self._bytes = peer_id_bytes
-
+
@staticmethod
def from_bytes(data: bytes) -> "PeerId":
"""Create PeerId from raw bytes."""
return PeerId(data)
-
+
@staticmethod
def from_public_key(public_key_bytes: bytes) -> "PeerId":
"""Create PeerId from a public key by hashing it."""
@@ -40,51 +40,51 @@ class PeerId:
# For larger keys, use SHA-256
hash_digest = hashlib.sha256(public_key_bytes).digest()
return PeerId(hash_digest)
-
+
def to_bytes(self) -> bytes:
"""Return the raw bytes of this PeerId."""
return self._bytes
-
+
def to_base58(self) -> str:
"""Return the base58-encoded string representation."""
- return base58.b58encode(self._bytes).decode('ascii')
-
+ return base58.b58encode(self._bytes).decode("ascii")
+
def __str__(self) -> str:
"""Return the base58-encoded string representation."""
return self.to_base58()
-
+
def __repr__(self) -> str:
"""Return debug representation."""
return f"PeerId('{self.to_base58()}')"
-
+
def __eq__(self, other: object) -> bool:
"""Check equality with another PeerId."""
if not isinstance(other, PeerId):
return False
return self._bytes == other._bytes
-
+
def __hash__(self) -> int:
"""Make PeerId hashable."""
return hash(self._bytes)
-@final
+@final
class Keypair:
"""
A py-libp2p compatible keypair implementation.
Provides the same interface as py-libp2p's KeyPair.
"""
-
+
def __init__(self, private_key: ed25519.Ed25519PrivateKey) -> None:
self._private_key = private_key
self._public_key = private_key.public_key()
-
+
@staticmethod
def generate_ed25519() -> "Keypair":
"""Generate a new Ed25519 keypair."""
private_key = ed25519.Ed25519PrivateKey.generate()
return Keypair(private_key)
-
+
@staticmethod
def from_protobuf_encoding(data: bytes) -> "Keypair":
"""
@@ -93,52 +93,54 @@ class Keypair:
"""
if len(data) < 2:
raise ValueError("Invalid protobuf data: too short")
-
+
# Simple protobuf parsing for our specific use case
# We expect: field 1 (type) as varint, field 2 (data) as bytes
offset = 0
-
+
# Parse type field (field tag 1, wire type 0 = varint)
if data[offset] != 0x08: # field 1, varint
raise ValueError("Expected type field")
offset += 1
-
+
key_type = data[offset]
offset += 1
-
+
if key_type != 1: # Ed25519
raise ValueError(f"Unsupported key type: {key_type}")
-
+
# Parse data field (field tag 2, wire type 2 = length-delimited)
if offset >= len(data) or data[offset] != 0x12: # field 2, bytes
raise ValueError("Expected data field")
offset += 1
-
+
# Parse length
data_length = data[offset]
offset += 1
-
+
if data_length not in (32, 64):
raise ValueError(f"Invalid Ed25519 private key length: {data_length}")
-
+
if offset + data_length > len(data):
raise ValueError("Truncated private key data")
-
- key_data = data[offset:offset + data_length]
-
+
+ key_data = data[offset : offset + data_length]
+
try:
if data_length == 64:
# libp2p format: 32 bytes private key seed + 32 bytes public key
private_key_seed = key_data[:32]
- private_key = ed25519.Ed25519PrivateKey.from_private_bytes(private_key_seed)
+ private_key = ed25519.Ed25519PrivateKey.from_private_bytes(
+ private_key_seed
+ )
else:
# Raw 32-byte private key
private_key = ed25519.Ed25519PrivateKey.from_private_bytes(key_data)
-
+
return Keypair(private_key)
except Exception as e:
raise ValueError(f"Invalid Ed25519 private key: {e}") from e
-
+
def to_protobuf_encoding(self) -> bytes:
"""
Serialize this keypair to libp2p protobuf encoding.
@@ -147,17 +149,16 @@ class Keypair:
private_key_bytes = self._private_key.private_bytes(
encoding=serialization.Encoding.Raw,
format=serialization.PrivateFormat.Raw,
- encryption_algorithm=serialization.NoEncryption()
+ encryption_algorithm=serialization.NoEncryption(),
)
-
+
public_key_bytes = self._public_key.public_bytes(
- encoding=serialization.Encoding.Raw,
- format=serialization.PublicFormat.Raw
+ encoding=serialization.Encoding.Raw, format=serialization.PublicFormat.Raw
)
-
+
# libp2p Ed25519 format: private key seed (32) + public key (32)
combined_key_data = private_key_bytes + public_key_bytes
-
+
# Build protobuf manually for our simple case
# Field 1 (type): tag=0x08, value=1 (Ed25519)
# Field 2 (data): tag=0x12, length=64, data=combined_key_data
@@ -165,21 +166,20 @@ class Keypair:
result.extend([0x08, 0x01]) # field 1: type = 1 (Ed25519)
result.extend([0x12, 0x40]) # field 2: length = 64 bytes
result.extend(combined_key_data)
-
+
return bytes(result)
-
+
def to_peer_id(self) -> PeerId:
"""Generate a PeerId from this keypair's public key."""
public_key_bytes = self._public_key.public_bytes(
- encoding=serialization.Encoding.Raw,
- format=serialization.PublicFormat.Raw
+ encoding=serialization.Encoding.Raw, format=serialization.PublicFormat.Raw
)
return PeerId.from_public_key(public_key_bytes)
-
+
def sign(self, data: bytes) -> bytes:
"""Sign data with this keypair's private key."""
return self._private_key.sign(data)
-
+
def verify(self, data: bytes, signature: bytes) -> bool:
"""Verify a signature against data using this keypair's public key."""
try:
@@ -187,22 +187,21 @@ class Keypair:
return True
except Exception:
return False
-
+
@property
def public_key_bytes(self) -> bytes:
"""Get the raw public key bytes."""
return self._public_key.public_bytes(
- encoding=serialization.Encoding.Raw,
- format=serialization.PublicFormat.Raw
+ encoding=serialization.Encoding.Raw, format=serialization.PublicFormat.Raw
)
-
- @property
+
+ @property
def private_key_bytes(self) -> bytes:
"""Get the raw private key bytes."""
return self._private_key.private_bytes(
encoding=serialization.Encoding.Raw,
format=serialization.PrivateFormat.Raw,
- encryption_algorithm=serialization.NoEncryption()
+ encryption_algorithm=serialization.NoEncryption(),
)
# py-libp2p compatibility properties
@@ -210,7 +209,7 @@ class Keypair:
def private_key(self) -> ed25519.Ed25519PrivateKey:
"""Access to the underlying private key for py-libp2p compatibility."""
return self._private_key
-
+
@property
def public_key(self) -> ed25519.Ed25519PublicKey:
"""Access to the underlying public key for py-libp2p compatibility."""
@@ -223,7 +222,9 @@ def ensure_type[T](obj: Any, expected_type: Type[T]) -> T: # type: ignore
return obj
-def get_node_id_keypair(path: str | bytes | os.PathLike[str] | os.PathLike[bytes] = EXO_NODE_ID_KEYPAIR) -> Keypair:
+def get_node_id_keypair(
+ path: str | bytes | os.PathLike[str] | os.PathLike[bytes] = EXO_NODE_ID_KEYPAIR,
+) -> Keypair:
"""
Obtains the :class:`Keypair` associated with this node-ID.
Obtain the :class:`PeerId` by from it.
@@ -234,7 +235,7 @@ def get_node_id_keypair(path: str | bytes | os.PathLike[str] | os.PathLike[bytes
# operate with cross-process lock to avoid race conditions
with FileLock(lock_path(path)):
- with open(path, 'a+b') as f: # opens in append-mode => starts at EOF
+ with open(path, "a+b") as f: # opens in append-mode => starts at EOF
# if non-zero EOF, then file exists => use to get node-ID
if f.tell() != 0:
f.seek(0) # go to start & read protobuf-encoded bytes
@@ -243,10 +244,12 @@ def get_node_id_keypair(path: str | bytes | os.PathLike[str] | os.PathLike[bytes
try: # if decoded successfully, save & return
return Keypair.from_protobuf_encoding(protobuf_encoded)
except ValueError as e: # on runtime error, assume corrupt file
- logging.warning(f"Encountered error when trying to get keypair: {e}")
+ logging.warning(
+ f"Encountered error when trying to get keypair: {e}"
+ )
# if no valid credentials, create new ones and persist
- with open(path, 'w+b') as f:
+ with open(path, "w+b") as f:
keypair = Keypair.generate_ed25519()
f.write(keypair.to_protobuf_encoding())
- return keypair
\ No newline at end of file
+ return keypair
diff --git a/src/exo/worker/common.py b/src/exo/worker/common.py
index 49c1e077..143061a7 100644
--- a/src/exo/worker/common.py
+++ b/src/exo/worker/common.py
@@ -27,7 +27,7 @@ class AssignedRunner(BaseModel):
runner: Optional[RunnerSupervisor] # set if the runner is 'up'
model_config = ConfigDict(arbitrary_types_allowed=True)
-
+
def status_update_event(self) -> RunnerStatusUpdated:
return RunnerStatusUpdated(
runner_id=self.runner_id,
diff --git a/src/exo/worker/download/conftest.py b/src/exo/worker/download/conftest.py
index eb96acd2..4cf8b936 100644
--- a/src/exo/worker/download/conftest.py
+++ b/src/exo/worker/download/conftest.py
@@ -7,7 +7,7 @@ from exo.shared.types.worker.shards import PipelineShardMetadata
@pytest.fixture
async def model_meta() -> ModelMetadata:
- return await get_model_meta('mlx-community/Llama-3.2-1B-Instruct-4bit')
+ return await get_model_meta("mlx-community/Llama-3.2-1B-Instruct-4bit")
@pytest.fixture
@@ -33,4 +33,4 @@ def pipeline_shard_meta(model_meta: ModelMetadata):
world_size=num_nodes,
)
- return _pipeline_shard_meta
\ No newline at end of file
+ return _pipeline_shard_meta
diff --git a/src/exo/worker/download/download_utils.py b/src/exo/worker/download/download_utils.py
index b88b7577..e2b4e8a2 100644
--- a/src/exo/worker/download/download_utils.py
+++ b/src/exo/worker/download/download_utils.py
@@ -17,25 +17,28 @@ from pydantic import BaseModel, DirectoryPath, Field, TypeAdapter
from exo.shared.constants import EXO_HOME
from exo.shared.types.worker.shards import ShardMetadata
from exo.worker.download.huggingface_utils import (
- filter_repo_objects,
- get_allow_patterns,
- get_auth_headers,
- get_hf_endpoint,
+ filter_repo_objects,
+ get_allow_patterns,
+ get_auth_headers,
+ get_hf_endpoint,
)
class ModelSafetensorsIndexMetadata(BaseModel):
- total_size: Annotated[int, Field(ge=0)]
+ total_size: Annotated[int, Field(ge=0)]
+
class ModelSafetensorsIndex(BaseModel):
- metadata: Optional[ModelSafetensorsIndexMetadata]
- weight_map: Dict[str, str]
+ metadata: Optional[ModelSafetensorsIndexMetadata]
+ weight_map: Dict[str, str]
+
class FileListEntry(BaseModel):
type: Literal["file", "directory"]
path: str
size: int | None = None
+
class RepoFileDownloadProgress(BaseModel):
"""Progress information for an individual file within a repository download."""
@@ -53,6 +56,7 @@ class RepoFileDownloadProgress(BaseModel):
class Config:
frozen = True
+
class RepoDownloadProgress(BaseModel):
"""Aggregated download progress information for a repository/shard combination.
@@ -87,352 +91,537 @@ class RepoDownloadProgress(BaseModel):
class Config:
frozen = True # allow use as dict keys if desired
+
def build_model_path(model_id: str) -> DirectoryPath:
- return EXO_HOME / "models" / model_id.replace("/", "--")
+ return EXO_HOME / "models" / model_id.replace("/", "--")
+
async def resolve_model_path_for_repo(repo_id: str) -> Path:
- return (await ensure_models_dir())/repo_id.replace("/", "--")
+ return (await ensure_models_dir()) / repo_id.replace("/", "--")
+
async def ensure_exo_home() -> Path:
- await aios.makedirs(EXO_HOME, exist_ok=True)
- return EXO_HOME
+ await aios.makedirs(EXO_HOME, exist_ok=True)
+ return EXO_HOME
+
async def has_exo_home_read_access() -> bool:
- try:
- return await aios.access(EXO_HOME, os.R_OK)
- except OSError:
- return False
+ try:
+ return await aios.access(EXO_HOME, os.R_OK)
+ except OSError:
+ return False
+
async def has_exo_home_write_access() -> bool:
- try:
- return await aios.access(EXO_HOME, os.W_OK)
- except OSError:
- return False
+ try:
+ return await aios.access(EXO_HOME, os.W_OK)
+ except OSError:
+ return False
+
async def ensure_models_dir() -> Path:
- models_dir = EXO_HOME/"models"
- await aios.makedirs(models_dir, exist_ok=True)
- return models_dir
+ models_dir = EXO_HOME / "models"
+ await aios.makedirs(models_dir, exist_ok=True)
+ return models_dir
+
async def delete_model(repo_id: str) -> bool:
- model_dir = await ensure_models_dir()/repo_id.replace("/", "--")
- if not await aios.path.exists(model_dir):
- return False
- await asyncio.to_thread(shutil.rmtree, model_dir, ignore_errors=False)
- return True
+ model_dir = await ensure_models_dir() / repo_id.replace("/", "--")
+ if not await aios.path.exists(model_dir):
+ return False
+ await asyncio.to_thread(shutil.rmtree, model_dir, ignore_errors=False)
+ return True
+
async def seed_models(seed_dir: Union[str, Path]):
- """Move model in resources folder of app to .cache/huggingface/hub"""
- source_dir = Path(seed_dir)
- dest_dir = await ensure_models_dir()
- for path in source_dir.iterdir():
- if path.is_dir() and path.name.startswith("models--"):
- dest_path = dest_dir/path.name
- if await aios.path.exists(dest_path):
- print('Skipping moving model to .cache directory')
- else:
+ """Move model in resources folder of app to .cache/huggingface/hub"""
+ source_dir = Path(seed_dir)
+ dest_dir = await ensure_models_dir()
+ for path in source_dir.iterdir():
+ if path.is_dir() and path.name.startswith("models--"):
+ dest_path = dest_dir / path.name
+ if await aios.path.exists(dest_path):
+ print("Skipping moving model to .cache directory")
+ else:
+ try:
+ await aios.rename(str(path), str(dest_path))
+ except Exception:
+ print(f"Error seeding model {path} to {dest_path}")
+ traceback.print_exc()
+
+
+async def fetch_file_list_with_cache(
+ repo_id: str, revision: str = "main", recursive: bool = False
+) -> List[FileListEntry]:
+ target_dir = (
+ (await ensure_models_dir()) / "caches" / str(repo_id).replace("/", "--")
+ )
+ await aios.makedirs(target_dir, exist_ok=True)
+ cache_file = (
+ target_dir / f"{repo_id.replace('/', '--')}--{revision}--file_list.json"
+ )
+ if await aios.path.exists(cache_file):
+ async with aiofiles.open(cache_file, "r") as f:
+ return TypeAdapter(List[FileListEntry]).validate_json(await f.read())
+ file_list = await fetch_file_list_with_retry(repo_id, revision, recursive=recursive)
+ await aios.makedirs(cache_file.parent, exist_ok=True)
+ async with aiofiles.open(cache_file, "w") as f:
+ await f.write(TypeAdapter(List[FileListEntry]).dump_json(file_list).decode())
+ return file_list
+
+
+async def fetch_file_list_with_retry(
+ repo_id: str, revision: str = "main", path: str = "", recursive: bool = False
+) -> List[FileListEntry]:
+ n_attempts = 30
+ for attempt in range(n_attempts):
try:
- await aios.rename(str(path), str(dest_path))
- except Exception:
- print(f"Error seeding model {path} to {dest_path}")
- traceback.print_exc()
-
-async def fetch_file_list_with_cache(repo_id: str, revision: str = "main", recursive: bool = False) -> List[FileListEntry]:
- target_dir = (await ensure_models_dir())/"caches"/str(repo_id).replace("/", "--")
- await aios.makedirs(target_dir, exist_ok=True)
- cache_file = target_dir/f"{repo_id.replace('/', '--')}--{revision}--file_list.json"
- if await aios.path.exists(cache_file):
- async with aiofiles.open(cache_file, 'r') as f:
- return TypeAdapter(List[FileListEntry]).validate_json(await f.read())
- file_list = await fetch_file_list_with_retry(repo_id, revision, recursive=recursive)
- await aios.makedirs(cache_file.parent, exist_ok=True)
- async with aiofiles.open(cache_file, 'w') as f:
- await f.write(TypeAdapter(List[FileListEntry]).dump_json(file_list).decode())
- return file_list
-
-
-async def fetch_file_list_with_retry(repo_id: str, revision: str = "main", path: str = "", recursive: bool = False) -> List[FileListEntry]:
- n_attempts = 30
- for attempt in range(n_attempts):
- try:
- return await _fetch_file_list(repo_id, revision, path, recursive)
- except Exception as e:
- if attempt == n_attempts - 1:
- raise e
- await asyncio.sleep(min(8, 0.1 * float(2.0 ** int(attempt))))
- raise Exception(f"Failed to fetch file list for {repo_id=} {revision=} {path=} {recursive=}")
-
-async def _fetch_file_list(repo_id: str, revision: str = "main", path: str = "", recursive: bool = False) -> List[FileListEntry]:
- api_url = f"{get_hf_endpoint()}/api/models/{repo_id}/tree/{revision}"
- url = f"{api_url}/{path}" if path else api_url
-
- headers = await get_auth_headers()
- async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=30, connect=10, sock_read=30, sock_connect=10)) as session, session.get(url, headers=headers) as response:
- if response.status == 200:
- data_json = await response.text()
- data = TypeAdapter(list[FileListEntry]).validate_json(data_json)
- files: list[FileListEntry] = []
- for item in data:
- if item.type == "file":
- files.append(FileListEntry.model_validate(item))
- elif item.type == "directory" and recursive:
- subfiles = await _fetch_file_list(repo_id, revision, item.path, recursive)
- files.extend(subfiles)
- return files
- else:
- raise Exception(f"Failed to fetch file list: {response.status}")
+ return await _fetch_file_list(repo_id, revision, path, recursive)
+ except Exception as e:
+ if attempt == n_attempts - 1:
+ raise e
+ await asyncio.sleep(min(8, 0.1 * float(2.0 ** int(attempt))))
+ raise Exception(
+ f"Failed to fetch file list for {repo_id=} {revision=} {path=} {recursive=}"
+ )
+
+
+async def _fetch_file_list(
+ repo_id: str, revision: str = "main", path: str = "", recursive: bool = False
+) -> List[FileListEntry]:
+ api_url = f"{get_hf_endpoint()}/api/models/{repo_id}/tree/{revision}"
+ url = f"{api_url}/{path}" if path else api_url
+
+ headers = await get_auth_headers()
+ async with (
+ aiohttp.ClientSession(
+ timeout=aiohttp.ClientTimeout(
+ total=30, connect=10, sock_read=30, sock_connect=10
+ )
+ ) as session,
+ session.get(url, headers=headers) as response,
+ ):
+ if response.status == 200:
+ data_json = await response.text()
+ data = TypeAdapter(list[FileListEntry]).validate_json(data_json)
+ files: list[FileListEntry] = []
+ for item in data:
+ if item.type == "file":
+ files.append(FileListEntry.model_validate(item))
+ elif item.type == "directory" and recursive:
+ subfiles = await _fetch_file_list(
+ repo_id, revision, item.path, recursive
+ )
+ files.extend(subfiles)
+ return files
+ else:
+ raise Exception(f"Failed to fetch file list: {response.status}")
+
async def calc_hash(path: Path, hash_type: Literal["sha1", "sha256"] = "sha1") -> str:
- hasher = hashlib.sha1() if hash_type == "sha1" else hashlib.sha256()
- if hash_type == "sha1":
- header = f"blob {(await aios.stat(path)).st_size}\0".encode()
- hasher.update(header)
- async with aiofiles.open(path, 'rb') as f:
- while chunk := await f.read(8 * 1024 * 1024):
- hasher.update(chunk)
- return hasher.hexdigest()
-
-
-async def file_meta(repo_id: str, revision: str, path: str, redirected_location: str | None = None) -> Tuple[int, str]:
- url = urljoin(f"{get_hf_endpoint()}/{repo_id}/resolve/{revision}/", path) if redirected_location is None else f"{get_hf_endpoint()}{redirected_location}"
+ hasher = hashlib.sha1() if hash_type == "sha1" else hashlib.sha256()
+ if hash_type == "sha1":
+ header = f"blob {(await aios.stat(path)).st_size}\0".encode()
+ hasher.update(header)
+ async with aiofiles.open(path, "rb") as f:
+ while chunk := await f.read(8 * 1024 * 1024):
+ hasher.update(chunk)
+ return hasher.hexdigest()
+
+
+async def file_meta(
+ repo_id: str, revision: str, path: str, redirected_location: str | None = None
+) -> Tuple[int, str]:
+ url = (
+ urljoin(f"{get_hf_endpoint()}/{repo_id}/resolve/{revision}/", path)
+ if redirected_location is None
+ else f"{get_hf_endpoint()}{redirected_location}"
+ )
headers = await get_auth_headers()
- async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=1800, connect=60, sock_read=1800, sock_connect=60)) as session, session.head(url, headers=headers) as r:
+ async with (
+ aiohttp.ClientSession(
+ timeout=aiohttp.ClientTimeout(
+ total=1800, connect=60, sock_read=1800, sock_connect=60
+ )
+ ) as session,
+ session.head(url, headers=headers) as r,
+ ):
if r.status == 307:
# Try to extract from X-Linked headers first (common for HF redirects)
- content_length = int(r.headers.get('x-linked-size') or r.headers.get('content-length') or 0)
- etag = r.headers.get('X-Linked-ETag') or r.headers.get('ETag') or r.headers.get('Etag')
+ content_length = int(
+ r.headers.get("x-linked-size") or r.headers.get("content-length") or 0
+ )
+ etag = (
+ r.headers.get("X-Linked-ETag")
+ or r.headers.get("ETag")
+ or r.headers.get("Etag")
+ )
if content_length > 0 and etag is not None:
- if (etag[0] == '"' and etag[-1] == '"') or (etag[0] == "'" and etag[-1] == "'"):
+ if (etag[0] == '"' and etag[-1] == '"') or (
+ etag[0] == "'" and etag[-1] == "'"
+ ):
etag = etag[1:-1]
return content_length, etag
# If not available, recurse with the redirect
- redirected_location = r.headers.get('Location')
+ redirected_location = r.headers.get("Location")
return await file_meta(repo_id, revision, path, redirected_location)
- content_length = int(r.headers.get('x-linked-size') or r.headers.get('content-length') or 0)
- etag = r.headers.get('X-Linked-ETag') or r.headers.get('ETag') or r.headers.get('Etag')
+ content_length = int(
+ r.headers.get("x-linked-size") or r.headers.get("content-length") or 0
+ )
+ etag = (
+ r.headers.get("X-Linked-ETag")
+ or r.headers.get("ETag")
+ or r.headers.get("Etag")
+ )
assert content_length > 0, f"No content length for {url}"
assert etag is not None, f"No remote hash for {url}"
if (etag[0] == '"' and etag[-1] == '"') or (etag[0] == "'" and etag[-1] == "'"):
etag = etag[1:-1]
return content_length, etag
-async def download_file_with_retry(repo_id: str, revision: str, path: str, target_dir: Path, on_progress: Callable[[int, int], None] = lambda _, __: None) -> Path:
- n_attempts = 30
- for attempt in range(n_attempts):
- try:
- return await _download_file(repo_id, revision, path, target_dir, on_progress)
- except Exception as e:
- if isinstance(e, FileNotFoundError) or attempt == n_attempts - 1:
- raise e
- print(f"Download error on attempt {attempt}/{n_attempts} for {repo_id=} {revision=} {path=} {target_dir=}")
- traceback.print_exc()
- await asyncio.sleep(min(8, 0.1 * (2.0 ** attempt)))
- raise Exception(f"Failed to download file {repo_id=} {revision=} {path=} {target_dir=}")
-
-async def _download_file(repo_id: str, revision: str, path: str, target_dir: Path, on_progress: Callable[[int, int], None] = lambda _, __: None) -> Path:
- if await aios.path.exists(target_dir/path):
- return target_dir/path
- await aios.makedirs((target_dir/path).parent, exist_ok=True)
- length, etag = await file_meta(repo_id, revision, path)
- remote_hash = etag[:-5] if etag.endswith("-gzip") else etag
- partial_path = target_dir/f"{path}.partial"
- resume_byte_pos = (await aios.stat(partial_path)).st_size if (await aios.path.exists(partial_path)) else None
- if resume_byte_pos != length:
- url = urljoin(f"{get_hf_endpoint()}/{repo_id}/resolve/{revision}/", path)
- headers = await get_auth_headers()
- if resume_byte_pos:
- headers['Range'] = f'bytes={resume_byte_pos}-'
- n_read = resume_byte_pos or 0
- async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=1800, connect=60, sock_read=1800, sock_connect=60)) as session, session.get(url, headers=headers, timeout=aiohttp.ClientTimeout(total=1800, connect=60, sock_read=1800, sock_connect=60)) as r:
- if r.status == 404:
- raise FileNotFoundError(f"File not found: {url}")
- assert r.status in [200, 206], f"Failed to download {path} from {url}: {r.status}"
- async with aiofiles.open(partial_path, 'ab' if resume_byte_pos else 'wb') as f:
- while chunk := await r.content.read(8 * 1024 * 1024):
- n_read = n_read + (await f.write(chunk))
- on_progress(n_read, length)
-
- final_hash = await calc_hash(partial_path, hash_type="sha256" if len(remote_hash) == 64 else "sha1")
- integrity = final_hash == remote_hash
- if not integrity:
- try:
- await aios.remove(partial_path)
- except Exception as e:
- print(f"Error removing partial file {partial_path}: {e}")
- raise Exception(f"Downloaded file {target_dir/path} has hash {final_hash} but remote hash is {remote_hash}")
- await aios.rename(partial_path, target_dir/path)
- return target_dir/path
-
-
-def calculate_repo_progress(shard: ShardMetadata, repo_id: str, revision: str, file_progress: Dict[str, RepoFileDownloadProgress], all_start_time: float) -> RepoDownloadProgress:
- all_total_bytes = sum(p.total for p in file_progress.values())
- all_downloaded_bytes = sum(p.downloaded for p in file_progress.values())
- all_downloaded_bytes_this_session = sum(p.downloaded_this_session for p in file_progress.values())
- elapsed_time = time.time() - all_start_time
- all_speed = all_downloaded_bytes_this_session / elapsed_time if elapsed_time > 0 else 0
- all_eta = timedelta(seconds=(all_total_bytes - all_downloaded_bytes) / all_speed) if all_speed > 0 else timedelta(seconds=0)
- status = (
- "complete"
- if all(p.status == "complete" for p in file_progress.values())
- else "in_progress" if any(p.status == "in_progress" for p in file_progress.values()) else "not_started"
- )
- return RepoDownloadProgress(
- repo_id=repo_id,
- repo_revision=revision,
- shard=shard,
- completed_files=len([p for p in file_progress.values() if p.downloaded == p.total]),
- total_files=len(file_progress),
- downloaded_bytes=all_downloaded_bytes,
- downloaded_bytes_this_session=all_downloaded_bytes_this_session,
- total_bytes=all_total_bytes,
- overall_speed=all_speed,
- overall_eta=all_eta,
- status=status,
- file_progress=file_progress,
- )
+
+async def download_file_with_retry(
+ repo_id: str,
+ revision: str,
+ path: str,
+ target_dir: Path,
+ on_progress: Callable[[int, int], None] = lambda _, __: None,
+) -> Path:
+ n_attempts = 30
+ for attempt in range(n_attempts):
+ try:
+ return await _download_file(
+ repo_id, revision, path, target_dir, on_progress
+ )
+ except Exception as e:
+ if isinstance(e, FileNotFoundError) or attempt == n_attempts - 1:
+ raise e
+ print(
+ f"Download error on attempt {attempt}/{n_attempts} for {repo_id=} {revision=} {path=} {target_dir=}"
+ )
+ traceback.print_exc()
+ await asyncio.sleep(min(8, 0.1 * (2.0**attempt)))
+ raise Exception(
+ f"Failed to download file {repo_id=} {revision=} {path=} {target_dir=}"
+ )
+
+
+async def _download_file(
+ repo_id: str,
+ revision: str,
+ path: str,
+ target_dir: Path,
+ on_progress: Callable[[int, int], None] = lambda _, __: None,
+) -> Path:
+ if await aios.path.exists(target_dir / path):
+ return target_dir / path
+ await aios.makedirs((target_dir / path).parent, exist_ok=True)
+ length, etag = await file_meta(repo_id, revision, path)
+ remote_hash = etag[:-5] if etag.endswith("-gzip") else etag
+ partial_path = target_dir / f"{path}.partial"
+ resume_byte_pos = (
+ (await aios.stat(partial_path)).st_size
+ if (await aios.path.exists(partial_path))
+ else None
+ )
+ if resume_byte_pos != length:
+ url = urljoin(f"{get_hf_endpoint()}/{repo_id}/resolve/{revision}/", path)
+ headers = await get_auth_headers()
+ if resume_byte_pos:
+ headers["Range"] = f"bytes={resume_byte_pos}-"
+ n_read = resume_byte_pos or 0
+ async with (
+ aiohttp.ClientSession(
+ timeout=aiohttp.ClientTimeout(
+ total=1800, connect=60, sock_read=1800, sock_connect=60
+ )
+ ) as session,
+ session.get(
+ url,
+ headers=headers,
+ timeout=aiohttp.ClientTimeout(
+ total=1800, connect=60, sock_read=1800, sock_connect=60
+ ),
+ ) as r,
+ ):
+ if r.status == 404:
+ raise FileNotFoundError(f"File not found: {url}")
+ assert r.status in [200, 206], (
+ f"Failed to download {path} from {url}: {r.status}"
+ )
+ async with aiofiles.open(
+ partial_path, "ab" if resume_byte_pos else "wb"
+ ) as f:
+ while chunk := await r.content.read(8 * 1024 * 1024):
+ n_read = n_read + (await f.write(chunk))
+ on_progress(n_read, length)
+
+ final_hash = await calc_hash(
+ partial_path, hash_type="sha256" if len(remote_hash) == 64 else "sha1"
+ )
+ integrity = final_hash == remote_hash
+ if not integrity:
+ try:
+ await aios.remove(partial_path)
+ except Exception as e:
+ print(f"Error removing partial file {partial_path}: {e}")
+ raise Exception(
+ f"Downloaded file {target_dir / path} has hash {final_hash} but remote hash is {remote_hash}"
+ )
+ await aios.rename(partial_path, target_dir / path)
+ return target_dir / path
+
+
+def calculate_repo_progress(
+ shard: ShardMetadata,
+ repo_id: str,
+ revision: str,
+ file_progress: Dict[str, RepoFileDownloadProgress],
+ all_start_time: float,
+) -> RepoDownloadProgress:
+ all_total_bytes = sum(p.total for p in file_progress.values())
+ all_downloaded_bytes = sum(p.downloaded for p in file_progress.values())
+ all_downloaded_bytes_this_session = sum(
+ p.downloaded_this_session for p in file_progress.values()
+ )
+ elapsed_time = time.time() - all_start_time
+ all_speed = (
+ all_downloaded_bytes_this_session / elapsed_time if elapsed_time > 0 else 0
+ )
+ all_eta = (
+ timedelta(seconds=(all_total_bytes - all_downloaded_bytes) / all_speed)
+ if all_speed > 0
+ else timedelta(seconds=0)
+ )
+ status = (
+ "complete"
+ if all(p.status == "complete" for p in file_progress.values())
+ else "in_progress"
+ if any(p.status == "in_progress" for p in file_progress.values())
+ else "not_started"
+ )
+ return RepoDownloadProgress(
+ repo_id=repo_id,
+ repo_revision=revision,
+ shard=shard,
+ completed_files=len(
+ [p for p in file_progress.values() if p.downloaded == p.total]
+ ),
+ total_files=len(file_progress),
+ downloaded_bytes=all_downloaded_bytes,
+ downloaded_bytes_this_session=all_downloaded_bytes_this_session,
+ total_bytes=all_total_bytes,
+ overall_speed=all_speed,
+ overall_eta=all_eta,
+ status=status,
+ file_progress=file_progress,
+ )
+
async def get_weight_map(repo_id: str, revision: str = "main") -> Dict[str, str]:
- target_dir = (await ensure_models_dir())/str(repo_id).replace("/", "--")
- await aios.makedirs(target_dir, exist_ok=True)
- index_file = await download_file_with_retry(repo_id, revision, "model.safetensors.index.json", target_dir)
- async with aiofiles.open(index_file, 'r') as f:
- index_data = ModelSafetensorsIndex.model_validate_json(await f.read())
- return index_data.weight_map
+ target_dir = (await ensure_models_dir()) / str(repo_id).replace("/", "--")
+ await aios.makedirs(target_dir, exist_ok=True)
+ index_file = await download_file_with_retry(
+ repo_id, revision, "model.safetensors.index.json", target_dir
+ )
+ async with aiofiles.open(index_file, "r") as f:
+ index_data = ModelSafetensorsIndex.model_validate_json(await f.read())
+ return index_data.weight_map
+
async def resolve_allow_patterns(shard: ShardMetadata) -> List[str]:
- try:
- weight_map = await get_weight_map(str(shard.model_meta.model_id))
- return get_allow_patterns(weight_map, shard)
- except Exception:
- print(f"Error getting weight map for {shard.model_meta.model_id=}")
- traceback.print_exc()
- return ["*"]
+ try:
+ weight_map = await get_weight_map(str(shard.model_meta.model_id))
+ return get_allow_patterns(weight_map, shard)
+ except Exception:
+ print(f"Error getting weight map for {shard.model_meta.model_id=}")
+ traceback.print_exc()
+ return ["*"]
+
async def get_downloaded_size(path: Path) -> int:
- partial_path = path.with_suffix(path.suffix + ".partial")
- if await aios.path.exists(path):
- return (await aios.stat(path)).st_size
- if await aios.path.exists(partial_path):
- return (await aios.stat(partial_path)).st_size
- return 0
-
-async def download_progress_for_local_path(repo_id: str, shard: ShardMetadata, local_path: Path) -> RepoDownloadProgress:
- # Scan local files for accurate progress reporting
- file_progress: Dict[str, RepoFileDownloadProgress] = {}
- total_files = 0
- total_bytes = 0
-
- if await aios.path.isdir(local_path):
- # Recursively count files and sizes
- for root, _, files in os.walk(local_path):
- for f in files:
- if f.endswith(('.safetensors', '.bin', '.pt', '.gguf', '.json')):
- file_path = Path(root) / f
- size = (await aios.stat(file_path)).st_size
- rel_path = str(file_path.relative_to(local_path))
- file_progress[rel_path] = RepoFileDownloadProgress(
- repo_id=repo_id,
- repo_revision="local",
- file_path=rel_path,
- downloaded=size,
+ partial_path = path.with_suffix(path.suffix + ".partial")
+ if await aios.path.exists(path):
+ return (await aios.stat(path)).st_size
+ if await aios.path.exists(partial_path):
+ return (await aios.stat(partial_path)).st_size
+ return 0
+
+
+async def download_progress_for_local_path(
+ repo_id: str, shard: ShardMetadata, local_path: Path
+) -> RepoDownloadProgress:
+ # Scan local files for accurate progress reporting
+ file_progress: Dict[str, RepoFileDownloadProgress] = {}
+ total_files = 0
+ total_bytes = 0
+
+ if await aios.path.isdir(local_path):
+ # Recursively count files and sizes
+ for root, _, files in os.walk(local_path):
+ for f in files:
+ if f.endswith((".safetensors", ".bin", ".pt", ".gguf", ".json")):
+ file_path = Path(root) / f
+ size = (await aios.stat(file_path)).st_size
+ rel_path = str(file_path.relative_to(local_path))
+ file_progress[rel_path] = RepoFileDownloadProgress(
+ repo_id=repo_id,
+ repo_revision="local",
+ file_path=rel_path,
+ downloaded=size,
+ downloaded_this_session=0,
+ total=size,
+ speed=0,
+ eta=timedelta(0),
+ status="complete",
+ start_time=time.time(),
+ )
+ total_files += 1
+ total_bytes += size
+ else:
+ raise ValueError(f"Local path {local_path} is not a directory")
+
+ return RepoDownloadProgress(
+ repo_id=repo_id,
+ repo_revision="local",
+ shard=shard,
+ completed_files=total_files,
+ total_files=total_files,
+ downloaded_bytes=total_bytes,
+ downloaded_bytes_this_session=0,
+ total_bytes=total_bytes,
+ overall_speed=0,
+ overall_eta=timedelta(0),
+ status="complete",
+ file_progress=file_progress,
+ )
+
+
+async def download_shard(
+ shard: ShardMetadata,
+ on_progress: Callable[[ShardMetadata, RepoDownloadProgress], None],
+ max_parallel_downloads: int = 8,
+ skip_download: bool = False,
+ allow_patterns: List[str] | None = None,
+) -> tuple[Path, RepoDownloadProgress]:
+ if not skip_download:
+ print(f"Downloading {shard.model_meta.model_id=}")
+
+ # Handle local paths
+ if await aios.path.exists(str(shard.model_meta.model_id)):
+ print(f"Using local model path {shard.model_meta.model_id}")
+ local_path = Path(str(shard.model_meta.model_id))
+ return local_path, await download_progress_for_local_path(
+ str(shard.model_meta.model_id), shard, local_path
+ )
+
+ revision = "main"
+ target_dir = await ensure_models_dir() / str(shard.model_meta.model_id).replace(
+ "/", "--"
+ )
+ if not skip_download:
+ await aios.makedirs(target_dir, exist_ok=True)
+
+ if not allow_patterns:
+ allow_patterns = await resolve_allow_patterns(shard)
+
+ print(f"Downloading {shard.model_meta.model_id=} with {allow_patterns=}")
+
+ all_start_time = time.time()
+ # TODO: currently not recursive. Some models might require subdirectories - thus this will need to be changed.
+ file_list = await fetch_file_list_with_cache(
+ str(shard.model_meta.model_id), revision, recursive=False
+ )
+ filtered_file_list = list(
+ filter_repo_objects(
+ file_list, allow_patterns=allow_patterns, key=lambda x: x.path
+ )
+ )
+ file_progress: Dict[str, RepoFileDownloadProgress] = {}
+
+ def on_progress_wrapper(file: FileListEntry, curr_bytes: int, total_bytes: int):
+ start_time = (
+ file_progress[file.path].start_time
+ if file.path in file_progress
+ else time.time()
+ )
+ downloaded_this_session = (
+ file_progress[file.path].downloaded_this_session
+ + (curr_bytes - file_progress[file.path].downloaded)
+ if file.path in file_progress
+ else curr_bytes
+ )
+ speed = (
+ downloaded_this_session / (time.time() - start_time)
+ if time.time() - start_time > 0
+ else 0
+ )
+ eta = (
+ timedelta(seconds=(total_bytes - curr_bytes) / speed)
+ if speed > 0
+ else timedelta(seconds=0)
+ )
+ file_progress[file.path] = RepoFileDownloadProgress(
+ repo_id=str(shard.model_meta.model_id),
+ repo_revision=revision,
+ file_path=file.path,
+ downloaded=curr_bytes,
+ downloaded_this_session=downloaded_this_session,
+ total=total_bytes,
+ speed=speed,
+ eta=eta,
+ status="complete" if curr_bytes == total_bytes else "in_progress",
+ start_time=start_time,
+ )
+ on_progress(
+ shard,
+ calculate_repo_progress(
+ shard,
+ str(shard.model_meta.model_id),
+ revision,
+ file_progress,
+ all_start_time,
+ ),
+ )
+
+ for file in filtered_file_list:
+ downloaded_bytes = await get_downloaded_size(target_dir / file.path)
+ file_progress[file.path] = RepoFileDownloadProgress(
+ repo_id=str(shard.model_meta.model_id),
+ repo_revision=revision,
+ file_path=file.path,
+ downloaded=downloaded_bytes,
downloaded_this_session=0,
- total=size,
+ total=file.size or 0,
speed=0,
eta=timedelta(0),
- status="complete",
- start_time=time.time()
- )
- total_files += 1
- total_bytes += size
- else:
- raise ValueError(f"Local path {local_path} is not a directory")
-
- return RepoDownloadProgress(
- repo_id=repo_id,
- repo_revision="local",
- shard=shard,
- completed_files=total_files,
- total_files=total_files,
- downloaded_bytes=total_bytes,
- downloaded_bytes_this_session=0,
- total_bytes=total_bytes,
- overall_speed=0,
- overall_eta=timedelta(0),
- status="complete",
- file_progress=file_progress,
- )
-
-async def download_shard(shard: ShardMetadata,
- on_progress: Callable[[ShardMetadata, RepoDownloadProgress], None],
- max_parallel_downloads: int = 8,
- skip_download: bool = False,
- allow_patterns: List[str] | None = None) -> tuple[Path, RepoDownloadProgress]:
- if not skip_download:
- print(f"Downloading {shard.model_meta.model_id=}")
-
- # Handle local paths
- if await aios.path.exists(str(shard.model_meta.model_id)):
- print(f"Using local model path {shard.model_meta.model_id}")
- local_path = Path(str(shard.model_meta.model_id))
- return local_path, await download_progress_for_local_path(str(shard.model_meta.model_id), shard, local_path)
-
- revision = "main"
- target_dir = await ensure_models_dir()/str(shard.model_meta.model_id).replace("/", "--")
- if not skip_download:
- await aios.makedirs(target_dir, exist_ok=True)
-
- if not allow_patterns:
- allow_patterns = await resolve_allow_patterns(shard)
-
- print(f"Downloading {shard.model_meta.model_id=} with {allow_patterns=}")
-
- all_start_time = time.time()
- # TODO: currently not recursive. Some models might require subdirectories - thus this will need to be changed.
- file_list = await fetch_file_list_with_cache(str(shard.model_meta.model_id), revision, recursive=False)
- filtered_file_list = list(filter_repo_objects(file_list, allow_patterns=allow_patterns, key=lambda x: x.path))
- file_progress: Dict[str, RepoFileDownloadProgress] = {}
- def on_progress_wrapper(file: FileListEntry, curr_bytes: int, total_bytes: int):
- start_time = file_progress[file.path].start_time if file.path in file_progress else time.time()
- downloaded_this_session = file_progress[file.path].downloaded_this_session + (curr_bytes - file_progress[file.path].downloaded) if file.path in file_progress else curr_bytes
- speed = downloaded_this_session / (time.time() - start_time) if time.time() - start_time > 0 else 0
- eta = timedelta(seconds=(total_bytes - curr_bytes) / speed) if speed > 0 else timedelta(seconds=0)
- file_progress[file.path] = RepoFileDownloadProgress(
- repo_id=str(shard.model_meta.model_id),
- repo_revision=revision,
- file_path=file.path,
- downloaded=curr_bytes,
- downloaded_this_session=downloaded_this_session,
- total=total_bytes,
- speed=speed,
- eta=eta,
- status="complete" if curr_bytes == total_bytes else "in_progress",
- start_time=start_time,
+ status="complete" if downloaded_bytes == file.size else "not_started",
+ start_time=time.time(),
+ )
+
+ semaphore = asyncio.Semaphore(max_parallel_downloads)
+
+ async def download_with_semaphore(file: FileListEntry):
+ async with semaphore:
+ await download_file_with_retry(
+ str(shard.model_meta.model_id),
+ revision,
+ file.path,
+ target_dir,
+ lambda curr_bytes, total_bytes: on_progress_wrapper(
+ file, curr_bytes, total_bytes
+ ),
+ )
+
+ if not skip_download:
+ await asyncio.gather(
+ *[download_with_semaphore(file) for file in filtered_file_list]
+ )
+ final_repo_progress = calculate_repo_progress(
+ shard, str(shard.model_meta.model_id), revision, file_progress, all_start_time
)
- on_progress(shard, calculate_repo_progress(shard, str(shard.model_meta.model_id), revision, file_progress, all_start_time))
- for file in filtered_file_list:
- downloaded_bytes = await get_downloaded_size(target_dir/file.path)
- file_progress[file.path] = RepoFileDownloadProgress(
- repo_id=str(shard.model_meta.model_id),
- repo_revision=revision,
- file_path=file.path,
- downloaded=downloaded_bytes,
- downloaded_this_session=0,
- total=file.size or 0,
- speed=0,
- eta=timedelta(0),
- status="complete" if downloaded_bytes == file.size else "not_started",
- start_time=time.time(),
- )
-
- semaphore = asyncio.Semaphore(max_parallel_downloads)
- async def download_with_semaphore(file: FileListEntry):
- async with semaphore:
- await download_file_with_retry(str(shard.model_meta.model_id), revision, file.path, target_dir, lambda curr_bytes, total_bytes: on_progress_wrapper(file, curr_bytes, total_bytes))
- if not skip_download:
- await asyncio.gather(*[download_with_semaphore(file) for file in filtered_file_list])
- final_repo_progress = calculate_repo_progress(shard, str(shard.model_meta.model_id), revision, file_progress, all_start_time)
- on_progress(shard, final_repo_progress)
- if gguf := next((f for f in filtered_file_list if f.path.endswith(".gguf")), None):
- return target_dir/gguf.path, final_repo_progress
- else:
- return target_dir, final_repo_progress
+ on_progress(shard, final_repo_progress)
+ if gguf := next((f for f in filtered_file_list if f.path.endswith(".gguf")), None):
+ return target_dir / gguf.path, final_repo_progress
+ else:
+ return target_dir, final_repo_progress
diff --git a/src/exo/worker/download/huggingface_utils.py b/src/exo/worker/download/huggingface_utils.py
index 56d118c5..837d5bc3 100644
--- a/src/exo/worker/download/huggingface_utils.py
+++ b/src/exo/worker/download/huggingface_utils.py
@@ -10,88 +10,107 @@ from exo.shared.types.worker.shards import ShardMetadata
T = TypeVar("T")
+
def filter_repo_objects(
- items: Iterable[T],
- *,
- allow_patterns: Optional[Union[List[str], str]] = None,
- ignore_patterns: Optional[Union[List[str], str]] = None,
- key: Optional[Callable[[T], str]] = None,
+ items: Iterable[T],
+ *,
+ allow_patterns: Optional[Union[List[str], str]] = None,
+ ignore_patterns: Optional[Union[List[str], str]] = None,
+ key: Optional[Callable[[T], str]] = None,
) -> Generator[T, None, None]:
- if isinstance(allow_patterns, str):
- allow_patterns = [allow_patterns]
- if isinstance(ignore_patterns, str):
- ignore_patterns = [ignore_patterns]
- if allow_patterns is not None:
- allow_patterns = [_add_wildcard_to_directories(p) for p in allow_patterns]
- if ignore_patterns is not None:
- ignore_patterns = [_add_wildcard_to_directories(p) for p in ignore_patterns]
-
- if key is None:
- def _identity(item: T) -> str:
- if isinstance(item, str):
- return item
- if isinstance(item, Path):
- return str(item)
- raise ValueError(f"Please provide `key` argument in `filter_repo_objects`: `{item}` is not a string.")
- key = _identity
-
- for item in items:
- path = key(item)
- if allow_patterns is not None and not any(fnmatch(path, r) for r in allow_patterns):
- continue
- if ignore_patterns is not None and any(fnmatch(path, r) for r in ignore_patterns):
- continue
- yield item
+ if isinstance(allow_patterns, str):
+ allow_patterns = [allow_patterns]
+ if isinstance(ignore_patterns, str):
+ ignore_patterns = [ignore_patterns]
+ if allow_patterns is not None:
+ allow_patterns = [_add_wildcard_to_directories(p) for p in allow_patterns]
+ if ignore_patterns is not None:
+ ignore_patterns = [_add_wildcard_to_directories(p) for p in ignore_patterns]
+
+ if key is None:
+
+ def _identity(item: T) -> str:
+ if isinstance(item, str):
+ return item
+ if isinstance(item, Path):
+ return str(item)
+ raise ValueError(
+ f"Please provide `key` argument in `filter_repo_objects`: `{item}` is not a string."
+ )
+
+ key = _identity
+
+ for item in items:
+ path = key(item)
+ if allow_patterns is not None and not any(
+ fnmatch(path, r) for r in allow_patterns
+ ):
+ continue
+ if ignore_patterns is not None and any(
+ fnmatch(path, r) for r in ignore_patterns
+ ):
+ continue
+ yield item
+
def _add_wildcard_to_directories(pattern: str) -> str:
- if pattern[-1] == "/":
- return pattern + "*"
- return pattern
+ if pattern[-1] == "/":
+ return pattern + "*"
+ return pattern
+
def get_hf_endpoint() -> str:
- return os.environ.get('HF_ENDPOINT', "https://huggingface.co")
+ return os.environ.get("HF_ENDPOINT", "https://huggingface.co")
+
def get_hf_home() -> Path:
- """Get the Hugging Face home directory."""
- return Path(os.environ.get("HF_HOME", Path.home()/".cache"/"huggingface"))
+ """Get the Hugging Face home directory."""
+ return Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface"))
+
async def get_hf_token() -> Optional[str]:
- """Retrieve the Hugging Face token from the user's HF_HOME directory."""
- token_path = get_hf_home()/"token"
- if await aios.path.exists(token_path):
- async with aiofiles.open(token_path, 'r') as f:
- return (await f.read()).strip()
- return None
+ """Retrieve the Hugging Face token from the user's HF_HOME directory."""
+ token_path = get_hf_home() / "token"
+ if await aios.path.exists(token_path):
+ async with aiofiles.open(token_path, "r") as f:
+ return (await f.read()).strip()
+ return None
+
async def get_auth_headers() -> dict[str, str]:
- """Get authentication headers if a token is available."""
- token = await get_hf_token()
- if token:
- return {"Authorization": f"Bearer {token}"}
- return {}
+ """Get authentication headers if a token is available."""
+ token = await get_hf_token()
+ if token:
+ return {"Authorization": f"Bearer {token}"}
+ return {}
+
def extract_layer_num(tensor_name: str) -> Optional[int]:
- # This is a simple example and might need to be adjusted based on the actual naming convention
- parts = tensor_name.split('.')
- for part in parts:
- if part.isdigit():
- return int(part)
- return None
+ # This is a simple example and might need to be adjusted based on the actual naming convention
+ parts = tensor_name.split(".")
+ for part in parts:
+ if part.isdigit():
+ return int(part)
+ return None
+
def get_allow_patterns(weight_map: Dict[str, str], shard: ShardMetadata) -> List[str]:
- default_patterns = set(["*.json", "*.py", "tokenizer.model", "*.tiktoken", "*.txt"])
- shard_specific_patterns: set[str] = set()
- if weight_map:
- for tensor_name, filename in weight_map.items():
- layer_num = extract_layer_num(tensor_name)
- if layer_num is not None and shard.start_layer <= layer_num <= shard.end_layer:
- shard_specific_patterns.add(filename)
- sorted_file_names = sorted(weight_map.values())
- if shard.is_first_layer:
- shard_specific_patterns.add(sorted_file_names[0])
- elif shard.is_last_layer:
- shard_specific_patterns.add(sorted_file_names[-1])
- else:
- shard_specific_patterns = set(["*.safetensors"])
- print(f"get_allow_patterns {shard=} {shard_specific_patterns=}")
- return list(default_patterns | shard_specific_patterns)
+ default_patterns = set(["*.json", "*.py", "tokenizer.model", "*.tiktoken", "*.txt"])
+ shard_specific_patterns: set[str] = set()
+ if weight_map:
+ for tensor_name, filename in weight_map.items():
+ layer_num = extract_layer_num(tensor_name)
+ if (
+ layer_num is not None
+ and shard.start_layer <= layer_num <= shard.end_layer
+ ):
+ shard_specific_patterns.add(filename)
+ sorted_file_names = sorted(weight_map.values())
+ if shard.is_first_layer:
+ shard_specific_patterns.add(sorted_file_names[0])
+ elif shard.is_last_layer:
+ shard_specific_patterns.add(sorted_file_names[-1])
+ else:
+ shard_specific_patterns = set(["*.safetensors"])
+ print(f"get_allow_patterns {shard=} {shard_specific_patterns=}")
+ return list(default_patterns | shard_specific_patterns)
diff --git a/src/exo/worker/download/impl_shard_downloader.py b/src/exo/worker/download/impl_shard_downloader.py
index 170e68c6..6f49c3fb 100644
--- a/src/exo/worker/download/impl_shard_downloader.py
+++ b/src/exo/worker/download/impl_shard_downloader.py
@@ -5,134 +5,185 @@ 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,
+ PartitionStrategy,
+ PipelineShardMetadata,
+ ShardMetadata,
)
from exo.worker.download.download_utils import RepoDownloadProgress, download_shard
from exo.worker.download.shard_downloader import ShardDownloader
def exo_shard_downloader(max_parallel_downloads: int = 8) -> ShardDownloader:
- return SingletonShardDownloader(CachedShardDownloader(ResumableShardDownloader(max_parallel_downloads)))
+ return SingletonShardDownloader(
+ CachedShardDownloader(ResumableShardDownloader(max_parallel_downloads))
+ )
-async def build_base_shard(model_id: str) -> Optional[ShardMetadata]:
- model_meta = await get_model_meta(model_id)
- # 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,
- end_layer=model_meta.n_layers,
- n_layers=model_meta.n_layers,
- )
-async def build_full_shard(model_id: str) -> Optional[PipelineShardMetadata]:
- base_shard = await build_base_shard(model_id)
- if base_shard is None:
- 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,
- end_layer=base_shard.n_layers,
- n_layers=base_shard.n_layers,
- )
+async def build_base_shard(model_id: str) -> Optional[ShardMetadata]:
+ model_meta = await get_model_meta(model_id)
+ # 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,
+ end_layer=model_meta.n_layers,
+ n_layers=model_meta.n_layers,
+ )
-class SingletonShardDownloader(ShardDownloader):
- def __init__(self, shard_downloader: ShardDownloader):
- self.shard_downloader = shard_downloader
- self.active_downloads: Dict[ShardMetadata, asyncio.Task[Path]] = {}
- def on_progress(self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]) -> None:
- self.shard_downloader.on_progress(callback)
+async def build_full_shard(model_id: str) -> Optional[PipelineShardMetadata]:
+ base_shard = await build_base_shard(model_id)
+ if base_shard is None:
+ 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,
+ end_layer=base_shard.n_layers,
+ n_layers=base_shard.n_layers,
+ )
- async def ensure_shard(self, shard: ShardMetadata, config_only: bool = False) -> Path:
- if shard not in self.active_downloads:
- self.active_downloads[shard] = asyncio.create_task(self.shard_downloader.ensure_shard(shard, config_only))
- try:
- return await self.active_downloads[shard]
- finally:
- if shard in self.active_downloads and self.active_downloads[shard].done():
- del self.active_downloads[shard]
- async def get_shard_download_status(self) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]:
- async for path, status in self.shard_downloader.get_shard_download_status():
- yield path, status
+class SingletonShardDownloader(ShardDownloader):
+ def __init__(self, shard_downloader: ShardDownloader):
+ self.shard_downloader = shard_downloader
+ self.active_downloads: Dict[ShardMetadata, asyncio.Task[Path]] = {}
+
+ def on_progress(
+ self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
+ ) -> None:
+ self.shard_downloader.on_progress(callback)
+
+ async def ensure_shard(
+ self, shard: ShardMetadata, config_only: bool = False
+ ) -> Path:
+ if shard not in self.active_downloads:
+ self.active_downloads[shard] = asyncio.create_task(
+ self.shard_downloader.ensure_shard(shard, config_only)
+ )
+ try:
+ return await self.active_downloads[shard]
+ finally:
+ if shard in self.active_downloads and self.active_downloads[shard].done():
+ del self.active_downloads[shard]
+
+ async def get_shard_download_status(
+ self,
+ ) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]:
+ async for path, status in self.shard_downloader.get_shard_download_status():
+ yield path, status
+
+ async def get_shard_download_status_for_shard(
+ self, shard: ShardMetadata
+ ) -> RepoDownloadProgress:
+ return await self.shard_downloader.get_shard_download_status_for_shard(shard)
- async def get_shard_download_status_for_shard(self, shard: ShardMetadata) -> RepoDownloadProgress:
- return await self.shard_downloader.get_shard_download_status_for_shard(shard)
class CachedShardDownloader(ShardDownloader):
- def __init__(self, shard_downloader: ShardDownloader):
- self.shard_downloader = shard_downloader
- self.cache: Dict[tuple[str, ShardMetadata], Path] = {}
-
- def on_progress(self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]) -> None:
- self.shard_downloader.on_progress(callback)
-
- async def ensure_shard(self, shard: ShardMetadata, config_only: bool = False) -> Path:
- if (shard.model_meta.model_id, shard) in self.cache:
- # print(f"ensure_shard cache hit {shard=}")
- return self.cache[(shard.model_meta.model_id, shard)]
+ def __init__(self, shard_downloader: ShardDownloader):
+ self.shard_downloader = shard_downloader
+ self.cache: Dict[tuple[str, ShardMetadata], Path] = {}
+
+ def on_progress(
+ self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
+ ) -> None:
+ self.shard_downloader.on_progress(callback)
+
+ async def ensure_shard(
+ self, shard: ShardMetadata, config_only: bool = False
+ ) -> Path:
+ if (shard.model_meta.model_id, shard) in self.cache:
+ # print(f"ensure_shard cache hit {shard=}")
+ return self.cache[(shard.model_meta.model_id, shard)]
+
+ # print(f"ensure_shard cache miss {shard=}")
+ target_dir = await self.shard_downloader.ensure_shard(shard, config_only)
+ self.cache[(shard.model_meta.model_id, shard)] = target_dir
+ return target_dir
+
+ async def get_shard_download_status(
+ self,
+ ) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]:
+ async for path, status in self.shard_downloader.get_shard_download_status():
+ yield path, status
+
+ async def get_shard_download_status_for_shard(
+ self, shard: ShardMetadata
+ ) -> RepoDownloadProgress:
+ return await self.shard_downloader.get_shard_download_status_for_shard(shard)
- # print(f"ensure_shard cache miss {shard=}")
- target_dir = await self.shard_downloader.ensure_shard(shard, config_only)
- self.cache[(shard.model_meta.model_id, shard)] = target_dir
- return target_dir
-
- async def get_shard_download_status(self) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]:
- async for path, status in self.shard_downloader.get_shard_download_status():
- yield path, status
-
- async def get_shard_download_status_for_shard(self, shard: ShardMetadata) -> RepoDownloadProgress:
- return await self.shard_downloader.get_shard_download_status_for_shard(shard)
class ResumableShardDownloader(ShardDownloader):
- def __init__(self, max_parallel_downloads: int = 8):
- self.max_parallel_downloads = max_parallel_downloads
- self.on_progress_callbacks: List[Callable[[ShardMetadata, RepoDownloadProgress], None]] = []
-
- def on_progress_wrapper(self, shard: ShardMetadata, progress: RepoDownloadProgress) -> None:
- for callback in self.on_progress_callbacks:
- callback(shard, progress)
-
- def on_progress(self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]) -> None:
- self.on_progress_callbacks.append(callback)
-
- async def ensure_shard(self, shard: ShardMetadata, config_only: bool = False) -> Path:
- allow_patterns = ["config.json"] if config_only else None
-
- # print(f"ensure_shard {shard=} {config_only=} {allow_patterns=}")
- target_dir, _ = await download_shard(shard, self.on_progress_wrapper, max_parallel_downloads=self.max_parallel_downloads, allow_patterns=allow_patterns)
- return target_dir
-
- async def get_shard_download_status(self) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]:
- # print("get_shard_download_status")
- async def _status_for_model(model_id: str) -> Optional[tuple[Path, RepoDownloadProgress]]:
- """Helper coroutine that builds the shard for a model and gets its download status."""
- shard = await build_full_shard(model_id)
- if shard is None:
- return None
- return await download_shard(shard, self.on_progress_wrapper, skip_download=True)
-
- # Kick off download status coroutines concurrently
- tasks = [asyncio.create_task(_status_for_model(model_card.model_id)) for model_card in MODEL_CARDS.values()]
-
- for task in asyncio.as_completed(tasks):
- try:
- result = await task
- if result is None:
- continue
- path, progress = result
- yield (path, progress)
- except Exception as e:
- print("Error downloading shard:", e)
-
- async def get_shard_download_status_for_shard(self, shard: ShardMetadata) -> RepoDownloadProgress:
- _, progress = await download_shard(shard, self.on_progress_wrapper, skip_download=True)
- return progress
+ def __init__(self, max_parallel_downloads: int = 8):
+ self.max_parallel_downloads = max_parallel_downloads
+ self.on_progress_callbacks: List[
+ Callable[[ShardMetadata, RepoDownloadProgress], None]
+ ] = []
+
+ def on_progress_wrapper(
+ self, shard: ShardMetadata, progress: RepoDownloadProgress
+ ) -> None:
+ for callback in self.on_progress_callbacks:
+ callback(shard, progress)
+
+ def on_progress(
+ self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
+ ) -> None:
+ self.on_progress_callbacks.append(callback)
+
+ async def ensure_shard(
+ self, shard: ShardMetadata, config_only: bool = False
+ ) -> Path:
+ allow_patterns = ["config.json"] if config_only else None
+
+ # print(f"ensure_shard {shard=} {config_only=} {allow_patterns=}")
+ target_dir, _ = await download_shard(
+ shard,
+ self.on_progress_wrapper,
+ max_parallel_downloads=self.max_parallel_downloads,
+ allow_patterns=allow_patterns,
+ )
+ return target_dir
+
+ async def get_shard_download_status(
+ self,
+ ) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]:
+ # print("get_shard_download_status")
+ async def _status_for_model(
+ model_id: str,
+ ) -> Optional[tuple[Path, RepoDownloadProgress]]:
+ """Helper coroutine that builds the shard for a model and gets its download status."""
+ shard = await build_full_shard(model_id)
+ if shard is None:
+ return None
+ return await download_shard(
+ shard, self.on_progress_wrapper, skip_download=True
+ )
+
+ # Kick off download status coroutines concurrently
+ tasks = [
+ asyncio.create_task(_status_for_model(model_card.model_id))
+ for model_card in MODEL_CARDS.values()
+ ]
+
+ for task in asyncio.as_completed(tasks):
+ try:
+ result = await task
+ if result is None:
+ continue
+ path, progress = result
+ yield (path, progress)
+ except Exception as e:
+ print("Error downloading shard:", e)
+
+ async def get_shard_download_status_for_shard(
+ self, shard: ShardMetadata
+ ) -> RepoDownloadProgress:
+ _, progress = await download_shard(
+ shard, self.on_progress_wrapper, skip_download=True
+ )
+ return progress
diff --git a/src/exo/worker/download/shard_downloader.py b/src/exo/worker/download/shard_downloader.py
index 6fcba625..ddb78915 100644
--- a/src/exo/worker/download/shard_downloader.py
+++ b/src/exo/worker/download/shard_downloader.py
@@ -5,18 +5,20 @@ from typing import AsyncIterator, Callable
from exo.shared.types.models import ModelMetadata
from exo.shared.types.worker.shards import (
- PartitionStrategy,
- PipelineShardMetadata,
- ShardMetadata,
+ PartitionStrategy,
+ PipelineShardMetadata,
+ ShardMetadata,
)
from exo.worker.download.download_utils import RepoDownloadProgress
# TODO: the PipelineShardMetadata getting reinstantiated is a bit messy. Shoudl this be a classmethod?
class ShardDownloader(ABC):
- @abstractmethod
- async def ensure_shard(self, shard: ShardMetadata, config_only: bool = False) -> Path:
- """
+ @abstractmethod
+ async def ensure_shard(
+ self, shard: ShardMetadata, config_only: bool = False
+ ) -> Path:
+ """
Ensures that the shard is downloaded.
Does not allow multiple overlapping downloads at once.
If you try to download a Shard which overlaps a Shard that is already being downloaded,
@@ -27,79 +29,108 @@ class ShardDownloader(ABC):
inference_engine_name (str): The inference engine used on the node hosting the shard
"""
- @abstractmethod
- def on_progress(self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]) -> None:
- pass
+ @abstractmethod
+ def on_progress(
+ self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
+ ) -> None:
+ pass
- @abstractmethod
- async def get_shard_download_status(self) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]:
- """Get the download status of shards.
+ @abstractmethod
+ async def get_shard_download_status(
+ self,
+ ) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]:
+ """Get the download status of shards.
- Yields:
- tuple[Path, RepoDownloadProgress]: The path and progress of a shard download.
- """
- yield (
- Path("/tmp/noop_shard"),
- RepoDownloadProgress(
- repo_id="noop",
- repo_revision="noop",
- shard=PipelineShardMetadata(
- model_meta=ModelMetadata(
- model_id='noop',
- pretty_name='noope',
- storage_size_kilobytes=0,
- n_layers=1
+ Yields:
+ tuple[Path, RepoDownloadProgress]: The path and progress of a shard download.
+ """
+ yield (
+ Path("/tmp/noop_shard"),
+ RepoDownloadProgress(
+ repo_id="noop",
+ repo_revision="noop",
+ shard=PipelineShardMetadata(
+ model_meta=ModelMetadata(
+ model_id="noop",
+ pretty_name="noope",
+ storage_size_kilobytes=0,
+ n_layers=1,
+ ),
+ partition_strategy=PartitionStrategy.pipeline,
+ device_rank=0,
+ world_size=1,
+ start_layer=0,
+ end_layer=1,
+ n_layers=1,
),
- partition_strategy=PartitionStrategy.pipeline,
- device_rank=0,
- world_size=1,
- start_layer=0,
- end_layer=1,
- n_layers=1,
+ completed_files=0,
+ total_files=0,
+ downloaded_bytes=0,
+ downloaded_bytes_this_session=0,
+ total_bytes=0,
+ overall_speed=0,
+ overall_eta=timedelta(seconds=0),
+ status="complete",
),
- completed_files=0,
- total_files=0,
- downloaded_bytes=0,
- downloaded_bytes_this_session=0,
- total_bytes=0,
- overall_speed=0,
- overall_eta=timedelta(seconds=0),
- status="complete",
)
- )
- @abstractmethod
- async def get_shard_download_status_for_shard(self, shard: ShardMetadata) -> RepoDownloadProgress:
- ...
+ @abstractmethod
+ async def get_shard_download_status_for_shard(
+ self, shard: ShardMetadata
+ ) -> RepoDownloadProgress: ...
class NoopShardDownloader(ShardDownloader):
- async def ensure_shard(self, shard: ShardMetadata, config_only: bool = False) -> Path:
- return Path("/tmp/noop_shard")
+ async def ensure_shard(
+ self, shard: ShardMetadata, config_only: bool = False
+ ) -> Path:
+ return Path("/tmp/noop_shard")
- def on_progress(self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]) -> None:
- pass
+ def on_progress(
+ self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
+ ) -> None:
+ pass
- async def get_shard_download_status(self) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]:
- yield (
- Path("/tmp/noop_shard"),
- RepoDownloadProgress(
- repo_id="noop",
- repo_revision="noop",
- shard=PipelineShardMetadata(
- model_meta=ModelMetadata(
- model_id='noop',
- pretty_name='noope',
- storage_size_kilobytes=0,
- n_layers=1
+ async def get_shard_download_status(
+ self,
+ ) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]:
+ yield (
+ Path("/tmp/noop_shard"),
+ RepoDownloadProgress(
+ repo_id="noop",
+ repo_revision="noop",
+ shard=PipelineShardMetadata(
+ model_meta=ModelMetadata(
+ model_id="noop",
+ pretty_name="noope",
+ storage_size_kilobytes=0,
+ n_layers=1,
+ ),
+ partition_strategy=PartitionStrategy.pipeline,
+ device_rank=0,
+ world_size=1,
+ start_layer=0,
+ end_layer=1,
+ n_layers=1,
),
- partition_strategy=PartitionStrategy.pipeline,
- device_rank=0,
- world_size=1,
- start_layer=0,
- end_layer=1,
- n_layers=1,
+ completed_files=0,
+ total_files=0,
+ downloaded_bytes=0,
+ downloaded_bytes_this_session=0,
+ total_bytes=0,
+ overall_speed=0,
+ overall_eta=timedelta(seconds=0),
+ status="complete",
),
+ )
+
+ async def get_shard_download_status_for_shard(
+ self, shard: ShardMetadata
+ ) -> RepoDownloadProgress:
+ return RepoDownloadProgress(
+ repo_id="noop",
+ repo_revision="noop",
+ shard=shard,
completed_files=0,
total_files=0,
downloaded_bytes=0,
@@ -109,19 +140,3 @@ class NoopShardDownloader(ShardDownloader):
overall_eta=timedelta(seconds=0),
status="complete",
)
- )
-
- async def get_shard_download_status_for_shard(self, shard: ShardMetadata) -> RepoDownloadProgress:
- return RepoDownloadProgress(
- repo_id="noop",
- repo_revision="noop",
- shard=shard,
- completed_files=0,
- total_files=0,
- downloaded_bytes=0,
- downloaded_bytes_this_session=0,
- total_bytes=0,
- overall_speed=0,
- overall_eta=timedelta(seconds=0),
- status="complete",
- )
\ No newline at end of file
diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py
index 189f668a..621d0cc1 100644
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -5,12 +5,12 @@ from exo.shared.apply import apply
from exo.shared.db.sqlite.event_log_manager import EventLogConfig, EventLogManager
from exo.shared.types.common import NodeId
from exo.shared.types.events import (
- NodePerformanceMeasured,
+ NodePerformanceMeasured,
)
from exo.shared.types.profiling import NodePerformanceProfile
from exo.shared.types.worker.ops import (
- ExecuteTaskOp,
- RunnerOp,
+ ExecuteTaskOp,
+ RunnerOp,
)
from exo.shared.utils import Keypair, get_node_id_keypair
from exo.worker.download.impl_shard_downloader import exo_shard_downloader
@@ -20,74 +20,94 @@ from exo.worker.worker import Worker
async def run(worker_state: Worker, logger: logging.Logger):
- assert worker_state.global_events is not None
-
- while True:
- # 1. get latest events
- events = await worker_state.global_events.get_events_since(worker_state.state.last_event_applied_idx)
-
- # 2. for each event, apply it to the state and run sagas
- for event_from_log in events:
- worker_state.state = apply(worker_state.state, event_from_log)
-
- # 3. based on the updated state, we plan & execute an operation.
- op: RunnerOp | None = plan(
- worker_state.assigned_runners,
- worker_state.node_id,
- worker_state.state.instances,
- worker_state.state.runners,
- worker_state.state.tasks,
- )
- if op is not None:
- worker_state.logger.info(f"!!! plan result: {op}")
-
- # run the op, synchronously blocking for now
- if op is not None:
- logger.info(f'Executing op {op}')
- try:
- async for event in worker_state.execute_op(op):
- await worker_state.event_publisher(event)
- except Exception as e:
- if isinstance(op, ExecuteTaskOp):
- generator = worker_state.fail_task(e, runner_id=op.runner_id, task_id=op.task.task_id)
- else:
- generator = worker_state.fail_runner(e, runner_id=op.runner_id)
-
- async for event in generator:
- await worker_state.event_publisher(event)
-
- await asyncio.sleep(0.01)
+ assert worker_state.global_events is not None
+ while True:
+ # 1. get latest events
+ events = await worker_state.global_events.get_events_since(
+ worker_state.state.last_event_applied_idx
+ )
+
+ # 2. for each event, apply it to the state and run sagas
+ for event_from_log in events:
+ worker_state.state = apply(worker_state.state, event_from_log)
+
+ # 3. based on the updated state, we plan & execute an operation.
+ op: RunnerOp | None = plan(
+ worker_state.assigned_runners,
+ worker_state.node_id,
+ worker_state.state.instances,
+ worker_state.state.runners,
+ worker_state.state.tasks,
+ )
+ if op is not None:
+ worker_state.logger.info(f"!!! plan result: {op}")
+
+ # run the op, synchronously blocking for now
+ if op is not None:
+ logger.info(f"Executing op {op}")
+ try:
+ async for event in worker_state.execute_op(op):
+ await worker_state.event_publisher(event)
+ except Exception as e:
+ if isinstance(op, ExecuteTaskOp):
+ generator = worker_state.fail_task(
+ e, runner_id=op.runner_id, task_id=op.task.task_id
+ )
+ else:
+ generator = worker_state.fail_runner(e, runner_id=op.runner_id)
+
+ async for event in generator:
+ await worker_state.event_publisher(event)
+ await asyncio.sleep(0.01)
async def async_main():
node_id_keypair: Keypair = get_node_id_keypair()
node_id = NodeId(node_id_keypair.to_peer_id().to_base58())
- logger: logging.Logger = logging.getLogger('worker_logger')
+ logger: logging.Logger = logging.getLogger("worker_logger")
logger.setLevel(logging.DEBUG)
if not logger.handlers:
handler = logging.StreamHandler()
- handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
+ handler.setFormatter(
+ logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
+ )
logger.addHandler(handler)
event_log_manager = EventLogManager(EventLogConfig(), logger)
await event_log_manager.initialize()
shard_downloader = exo_shard_downloader()
-
+
# TODO: add profiling etc to resource monitor
- async def resource_monitor_callback(node_performance_profile: NodePerformanceProfile) -> None:
+ async def resource_monitor_callback(
+ node_performance_profile: NodePerformanceProfile,
+ ) -> None:
await event_log_manager.worker_events.append_events(
- [NodePerformanceMeasured(node_id=node_id, node_profile=node_performance_profile)], origin=node_id
+ [
+ NodePerformanceMeasured(
+ node_id=node_id, node_profile=node_performance_profile
+ )
+ ],
+ origin=node_id,
)
+
asyncio.create_task(start_polling_node_metrics(callback=resource_monitor_callback))
- worker = Worker(node_id, logger, shard_downloader, event_log_manager.worker_events, event_log_manager.global_events)
+ worker = Worker(
+ node_id,
+ logger,
+ shard_downloader,
+ event_log_manager.worker_events,
+ event_log_manager.global_events,
+ )
await run(worker, logger)
+
def main():
asyncio.run(async_main())
+
if __name__ == "__main__":
main()
diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py
index a0be6920..1e97e1cf 100644
--- a/src/exo/worker/plan.py
+++ b/src/exo/worker/plan.py
@@ -28,7 +28,11 @@ from exo.shared.types.worker.runners import (
from exo.worker.common import AssignedRunner
-def unassign_runners(instances: Mapping[InstanceId, Instance], state_runners: Mapping[RunnerId, RunnerStatus], assigned_runners: dict[RunnerId, AssignedRunner]) -> UnassignRunnerOp | None:
+def unassign_runners(
+ instances: Mapping[InstanceId, Instance],
+ state_runners: Mapping[RunnerId, RunnerStatus],
+ assigned_runners: dict[RunnerId, AssignedRunner],
+) -> UnassignRunnerOp | None:
runner_ids: set[RunnerId] = {
runner_id
for instance in instances.values()
@@ -40,77 +44,122 @@ def unassign_runners(instances: Mapping[InstanceId, Instance], state_runners: Ma
# 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):
+ if assigned_runner_id in state_runners and isinstance(
+ state_runners[assigned_runner_id], DownloadingRunnerStatus
+ ):
return UnassignRunnerOp(runner_id=assigned_runner_id)
return None
-def failed_runners(assigned_runners: dict[RunnerId, AssignedRunner]) -> RunnerFailedOp | None:
+
+def failed_runners(
+ assigned_runners: dict[RunnerId, AssignedRunner],
+) -> RunnerFailedOp | None:
for runner_id, assigned_runner in assigned_runners.items():
- if assigned_runner.runner is not None and \
- not assigned_runner.runner.healthy and \
- not isinstance(assigned_runner.status, FailedRunnerStatus):
+ if (
+ assigned_runner.runner is not None
+ and not assigned_runner.runner.healthy
+ and not isinstance(assigned_runner.status, FailedRunnerStatus)
+ ):
return RunnerFailedOp(runner_id=runner_id)
return None
+
def spin_down_runners(
- instances: Mapping[InstanceId, Instance],
- assigned_runners: dict[RunnerId, AssignedRunner],
+ instances: Mapping[InstanceId, Instance],
+ assigned_runners: dict[RunnerId, AssignedRunner],
state_runners: Mapping[RunnerId, RunnerStatus],
- worker_node_id: NodeId) -> RunnerDownOp | None:
+ worker_node_id: NodeId,
+) -> RunnerDownOp | None:
for _instance_id, instance in instances.items():
for node_id, runner_id in instance.shard_assignments.node_to_runner.items():
if node_id != worker_node_id:
continue
# We spin down a runner if it's meant to be inactive and it's Loaded.
- if runner_id in assigned_runners and \
- isinstance(assigned_runners[runner_id].status, LoadedRunnerStatus) and \
- instance.instance_type == InstanceStatus.INACTIVE:
+ if (
+ runner_id in assigned_runners
+ and isinstance(assigned_runners[runner_id].status, LoadedRunnerStatus)
+ and instance.instance_type == InstanceStatus.INACTIVE
+ ):
return RunnerDownOp(runner_id=runner_id)
# If we are part of an instance that has a dead node - and we aren't the dead node - we should spin down
for _instance_id, instance in instances.items():
- if worker_node_id in instance.shard_assignments.node_to_runner and \
- instance.shard_assignments.node_to_runner[worker_node_id] in assigned_runners and \
- not isinstance(assigned_runners[instance.shard_assignments.node_to_runner[worker_node_id]].status, InactiveRunnerStatus): # make sure that our runner has not already been spun down into ready state
+ if (
+ worker_node_id in instance.shard_assignments.node_to_runner
+ and instance.shard_assignments.node_to_runner[worker_node_id]
+ in assigned_runners
+ and not isinstance(
+ assigned_runners[
+ instance.shard_assignments.node_to_runner[worker_node_id]
+ ].status,
+ InactiveRunnerStatus,
+ )
+ ): # make sure that our runner has not already been spun down into ready state
other_node_in_instance_has_failed = False
for runner_id in instance.shard_assignments.runner_to_shard:
- if runner_id in state_runners and \
- isinstance(state_runners[runner_id], FailedRunnerStatus) and \
- runner_id not in assigned_runners:
- other_node_in_instance_has_failed= True
+ if (
+ runner_id in state_runners
+ and isinstance(state_runners[runner_id], FailedRunnerStatus)
+ and runner_id not in assigned_runners
+ ):
+ other_node_in_instance_has_failed = True
if other_node_in_instance_has_failed:
# Spin down *our* runner
- return RunnerDownOp(runner_id=instance.shard_assignments.node_to_runner[worker_node_id])
+ return RunnerDownOp(
+ runner_id=instance.shard_assignments.node_to_runner[worker_node_id]
+ )
# If we are failed - and *all of the other nodes have spun down* - then we can spin down too.
for _instance_id, instance in instances.items():
- if worker_node_id in instance.shard_assignments.node_to_runner and \
- instance.shard_assignments.node_to_runner[worker_node_id] in state_runners and \
- instance.shard_assignments.node_to_runner[worker_node_id] in assigned_runners and \
- isinstance(assigned_runners[instance.shard_assignments.node_to_runner[worker_node_id]].status, FailedRunnerStatus):
-
+ if (
+ worker_node_id in instance.shard_assignments.node_to_runner
+ and instance.shard_assignments.node_to_runner[worker_node_id]
+ in state_runners
+ and instance.shard_assignments.node_to_runner[worker_node_id]
+ in assigned_runners
+ and isinstance(
+ assigned_runners[
+ instance.shard_assignments.node_to_runner[worker_node_id]
+ ].status,
+ FailedRunnerStatus,
+ )
+ ):
num_spundown_nodes = 0
for runner_id in instance.shard_assignments.runner_to_shard:
- if runner_id in state_runners and \
- isinstance(state_runners[runner_id], InactiveRunnerStatus) and \
- runner_id not in assigned_runners:
+ if (
+ runner_id in state_runners
+ and isinstance(state_runners[runner_id], InactiveRunnerStatus)
+ and runner_id not in assigned_runners
+ ):
num_spundown_nodes += 1
# Suggested:
# if runner_id in state_runners and isinstance(state.runners[runner_id], InactiveRunnerStatus):
# if runner_id != instance.shard_assignments.node_to_runner[worker_node_id]:
# num_spundown_nodes += 1
- if num_spundown_nodes == next(iter(instance.shard_assignments.runner_to_shard.values())).world_size - 1:
+ if (
+ num_spundown_nodes
+ == next(
+ iter(instance.shard_assignments.runner_to_shard.values())
+ ).world_size
+ - 1
+ ):
# All the other nodes are spun down - so now we can spin down too.
# This also catches the case of 1-node. If there's one node in the instance then we should spin down straight away
- return RunnerDownOp(runner_id=instance.shard_assignments.node_to_runner[worker_node_id])
+ return RunnerDownOp(
+ runner_id=instance.shard_assignments.node_to_runner[worker_node_id]
+ )
return None
-def assign_runners(instances: Mapping[InstanceId, Instance], assigned_runners: dict[RunnerId, AssignedRunner], worker_node_id: NodeId) -> AssignRunnerOp | None:
+
+def assign_runners(
+ instances: Mapping[InstanceId, Instance],
+ assigned_runners: dict[RunnerId, AssignedRunner],
+ worker_node_id: NodeId,
+) -> AssignRunnerOp | None:
for instance_id, instance in instances.items():
for node_id, runner_id in instance.shard_assignments.node_to_runner.items():
if node_id != worker_node_id:
@@ -120,29 +169,54 @@ def assign_runners(instances: Mapping[InstanceId, Instance], assigned_runners: d
return AssignRunnerOp(
runner_id=runner_id,
instance_id=instance_id,
- shard_metadata=instance.shard_assignments.runner_to_shard[runner_id],
- hosts=instance.hosts
+ shard_metadata=instance.shard_assignments.runner_to_shard[
+ runner_id
+ ],
+ hosts=instance.hosts,
)
return None
-def spin_up_runners(instances: Mapping[InstanceId, Instance], assigned_runners: dict[RunnerId, AssignedRunner], state_runners: Mapping[RunnerId, RunnerStatus], worker_node_id: NodeId) -> RunnerUpOp | None:
- for _instance_id, instance in instances.items():
- if worker_node_id in instance.shard_assignments.node_to_runner and \
- assigned_runners[instance.shard_assignments.node_to_runner[worker_node_id]].runner is None and \
- instance.instance_type == InstanceStatus.ACTIVE:
+def spin_up_runners(
+ instances: Mapping[InstanceId, Instance],
+ assigned_runners: dict[RunnerId, AssignedRunner],
+ state_runners: Mapping[RunnerId, RunnerStatus],
+ worker_node_id: NodeId,
+) -> RunnerUpOp | None:
+ for _instance_id, instance in instances.items():
+ if (
+ worker_node_id in instance.shard_assignments.node_to_runner
+ and assigned_runners[
+ instance.shard_assignments.node_to_runner[worker_node_id]
+ ].runner
+ is None
+ 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 != RunnerStatusType.Inactive:
+ if (
+ runner_id in state_runners
+ and state_runners[runner_id].runner_status
+ != RunnerStatusType.Inactive
+ ):
ready_to_spin = False
if ready_to_spin:
- return RunnerUpOp(runner_id=instance.shard_assignments.node_to_runner[worker_node_id])
+ return RunnerUpOp(
+ runner_id=instance.shard_assignments.node_to_runner[worker_node_id]
+ )
return None
-def execute_task_op(instances: Mapping[InstanceId, Instance], assigned_runners: dict[RunnerId, AssignedRunner], state_runners: Mapping[RunnerId, RunnerStatus], tasks: Mapping[TaskId, Task], worker_node_id: NodeId) -> ExecuteTaskOp | None:
+
+def execute_task_op(
+ instances: Mapping[InstanceId, Instance],
+ assigned_runners: dict[RunnerId, AssignedRunner],
+ state_runners: Mapping[RunnerId, RunnerStatus],
+ tasks: Mapping[TaskId, Task],
+ worker_node_id: NodeId,
+) -> ExecuteTaskOp | None:
for instance_id, instance in instances.items():
for node_id, runner_id in instance.shard_assignments.node_to_runner.items():
if node_id != worker_node_id:
@@ -150,21 +224,31 @@ def execute_task_op(instances: Mapping[InstanceId, Instance], assigned_runners:
assert runner_id in assigned_runners
runner = assigned_runners[runner_id]
if runner.status.runner_status != RunnerStatusType.Loaded:
- continue # The only previous state to get to Running is from Loaded
+ 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 == TaskStatus.PENDING
+ or task.task_status == TaskStatus.FAILED
):
- if (runner.shard_metadata.device_rank >= 1 or runner.shard_metadata.world_size == 1):
+ if (
+ runner.shard_metadata.device_rank >= 1
+ or runner.shard_metadata.world_size == 1
+ ):
return ExecuteTaskOp(runner_id=runner_id, task=task)
else:
# We already know our own status is Loaded. We are rank 0,
# so let's check that all the other runners are running - ready for us to fire the prompt.
running_runner_count = 0
- for other_runner_id, other_runner_status in state_runners.items():
- if other_runner_id in instance.shard_assignments.node_to_runner.values() and \
- isinstance(other_runner_status, RunningRunnerStatus):
+ for (
+ other_runner_id,
+ other_runner_status,
+ ) in state_runners.items():
+ if (
+ other_runner_id
+ in instance.shard_assignments.node_to_runner.values()
+ and isinstance(other_runner_status, RunningRunnerStatus)
+ ):
running_runner_count += 1
if running_runner_count == runner.shard_metadata.world_size - 1:
@@ -173,12 +257,13 @@ def execute_task_op(instances: Mapping[InstanceId, Instance], assigned_runners:
return None
-
-def plan(assigned_runners: dict[RunnerId, AssignedRunner],
- worker_node_id: NodeId,
- instances: Mapping[InstanceId, Instance],
- state_runners: Mapping[RunnerId, RunnerStatus], # all global
- tasks: Mapping[TaskId, Task]) -> RunnerOp | None:
+def plan(
+ assigned_runners: dict[RunnerId, AssignedRunner],
+ worker_node_id: NodeId,
+ instances: Mapping[InstanceId, Instance],
+ state_runners: Mapping[RunnerId, RunnerStatus], # all global
+ tasks: Mapping[TaskId, Task],
+) -> RunnerOp | None:
# First, unassign assigned runners that are no longer in the state.
if unop := unassign_runners(instances, state_runners, assigned_runners):
return unop
@@ -188,7 +273,9 @@ def plan(assigned_runners: dict[RunnerId, AssignedRunner],
return failed_op
# spin down runners that are no longer needed
- if down_op := spin_down_runners(instances, assigned_runners, state_runners, worker_node_id):
+ if down_op := spin_down_runners(
+ instances, assigned_runners, state_runners, worker_node_id
+ ):
return down_op
# Then assign runners we do want
@@ -196,11 +283,15 @@ def plan(assigned_runners: dict[RunnerId, AssignedRunner],
return assign_op
# Then spin up 'ready' runners that should be active
- if runner_up_op := spin_up_runners(instances, assigned_runners, state_runners, worker_node_id):
+ if runner_up_op := spin_up_runners(
+ instances, assigned_runners, state_runners, worker_node_id
+ ):
return runner_up_op
# Then make sure things are running based on tasks.
- if exec_op := execute_task_op(instances, assigned_runners, state_runners, tasks, worker_node_id):
+ if exec_op := execute_task_op(
+ instances, assigned_runners, state_runners, tasks, worker_node_id
+ ):
return exec_op
return None
diff --git a/src/exo/worker/runner/communication.py b/src/exo/worker/runner/communication.py
index 044d10c5..544bf4e8 100644
--- a/src/exo/worker/runner/communication.py
+++ b/src/exo/worker/runner/communication.py
@@ -31,7 +31,7 @@ async def runner_read_message() -> RunnerMessage:
loop = asyncio.get_running_loop()
line: bytes = await loop.run_in_executor(None, sys.stdin.buffer.readline)
- if not line: # This seems to be what triggers when we don't clean up the runner neatly and leave the process dangling.
+ if not line: # This seems to be what triggers when we don't clean up the runner neatly and leave the process dangling.
raise EOFError("No more data to read when reading runner message")
line = line.strip()
@@ -88,11 +88,11 @@ def runner_write_error(error: Exception) -> None:
# Skip writing error if it's a BrokenPipeError - supervisor is already gone
if isinstance(error, BrokenPipeError):
sys.exit(0)
-
+
error_response: ErrorResponse = ErrorResponse(
type=RunnerResponseType.ErrorResponse,
error_type=type(error).__name__,
error_message=str(error),
traceback=traceback.format_exc(),
)
- runner_write_response(error_response)
\ No newline at end of file
+ runner_write_response(error_response)
diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py
index fef5645c..25e1a025 100644
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -93,7 +93,6 @@ async def _mlx_generate(
if isinstance(item, Exception):
raise item
-
assert isinstance(item, GenerationResponse) # constrain datatype
runner_print(item.text)
yield item
@@ -113,10 +112,10 @@ async def main():
# For testing - these are fake break conditions
if model_shard_meta.immediate_exception:
- raise Exception('Fake exception - runner failed to spin up.')
+ raise Exception("Fake exception - runner failed to spin up.")
if model_shard_meta.should_timeout:
await asyncio.sleep(model_shard_meta.should_timeout)
-
+
setup_start_time = time.time()
mlx_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
@@ -133,8 +132,10 @@ async def main():
tokenizer=tokenizer,
sampler=sampler,
)
- runner_print(f'Warmed up by generating {toks} tokens')
- runner_write_response(InitializedResponse(time_taken=time.time() - setup_start_time))
+ runner_print(f"Warmed up by generating {toks} tokens")
+ runner_write_response(
+ InitializedResponse(time_taken=time.time() - setup_start_time)
+ )
while True:
message: RunnerMessage = await runner_read_message()
@@ -144,12 +145,23 @@ async def main():
# Ensure we have a chat-completion task subtype
# TODO: this is a hack, why are we only looking at the first message? should have a tokenizer
prompt = task.messages[0]
- if prompt.content is not None and 'EXO RUNNER MUST FAIL' in prompt.content:
- runner_print('raising exception')
- raise Exception('Artificial runner exception - for testing purposes only.')
- if prompt.content is not None and 'EXO RUNNER MUST OOM' in prompt.content:
+ if (
+ prompt.content is not None
+ and "EXO RUNNER MUST FAIL" in prompt.content
+ ):
+ runner_print("raising exception")
+ raise Exception(
+ "Artificial runner exception - for testing purposes only."
+ )
+ if (
+ prompt.content is not None
+ and "EXO RUNNER MUST OOM" in prompt.content
+ ):
mlx_force_oom()
- if prompt.content is not None and 'EXO RUNNER MUST TIMEOUT' in prompt.content:
+ if (
+ prompt.content is not None
+ and "EXO RUNNER MUST TIMEOUT" in prompt.content
+ ):
await asyncio.sleep(100)
# Generate responses using the actual MLX generation
diff --git a/src/exo/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py
index be94d2cc..bb9106d9 100644
--- a/src/exo/worker/runner/runner_supervisor.py
+++ b/src/exo/worker/runner/runner_supervisor.py
@@ -76,13 +76,11 @@ class RunnerSupervisor:
The .create() classmethod pattern is used to ensure the constructor is asynchronous.
"""
cmd: list[str] = get_runner_command()
- runner_process = (
- await asyncio.create_subprocess_exec(
- *cmd,
- stdin=asyncio.subprocess.PIPE,
- stdout=asyncio.subprocess.PIPE,
- stderr=asyncio.subprocess.PIPE,
- )
+ runner_process = await asyncio.create_subprocess_exec(
+ *cmd,
+ stdin=asyncio.subprocess.PIPE,
+ stdout=asyncio.subprocess.PIPE,
+ stderr=asyncio.subprocess.PIPE,
)
read_queue: asyncio.Queue[RunnerResponse] = asyncio.Queue()
@@ -99,11 +97,13 @@ class RunnerSupervisor:
stderr_queue=stderr_queue,
)
- self.logger.info(f'initializing mlx instance with {model_shard_meta=}')
- await self.write_queue.put(SetupMessage(
- model_shard_meta=model_shard_meta,
- hosts=hosts,
- ))
+ self.logger.info(f"initializing mlx instance with {model_shard_meta=}")
+ await self.write_queue.put(
+ SetupMessage(
+ model_shard_meta=model_shard_meta,
+ hosts=hosts,
+ )
+ )
if not initialize_timeout:
initialize_timeout = get_init_timeout(model_shard_meta)
@@ -111,37 +111,42 @@ class RunnerSupervisor:
response = await self._read_with_error_check(initialize_timeout)
assert isinstance(response, InitializedResponse)
- self.logger.info(f'Runner initialized in {response.time_taken} seconds')
+ self.logger.info(f"Runner initialized in {response.time_taken} seconds")
return self
-
async def _read_with_error_check(self, timeout: float) -> RunnerResponse:
"""
Read from the queue with a timeout, but also check if the read_task has failed.
"""
queue_task = asyncio.create_task(self.read_queue.get())
-
+
done, pending = await asyncio.wait(
[queue_task, self.read_task],
timeout=timeout,
- return_when=asyncio.FIRST_COMPLETED
+ return_when=asyncio.FIRST_COMPLETED,
)
-
+
for task in pending:
if task is queue_task:
task.cancel()
-
+
if queue_task in done:
response = await queue_task
if isinstance(response, ErrorResponse):
- raise RunnerError(response.error_type, response.error_message, response.traceback or "")
+ raise RunnerError(
+ response.error_type,
+ response.error_message,
+ response.traceback or "",
+ )
return response
-
+
if self.read_task in done:
await self.read_task # Re-raises any exception from read_task
- self.logger.error('Unreachable code run. We should have raised an error on the read_task being done.')
-
+ self.logger.error(
+ "Unreachable code run. We should have raised an error on the read_task being done."
+ )
+
# if we haven't read from the queue, we have timed out.
await self.astop()
raise asyncio.TimeoutError()
@@ -149,7 +154,8 @@ class RunnerSupervisor:
async def stream_response(
self,
task: Task,
- request_started_callback: Callable[..., CoroutineType[Any, Any, None]] | None = None,
+ request_started_callback: Callable[..., CoroutineType[Any, Any, None]]
+ | None = None,
) -> AsyncGenerator[GenerationChunk]:
"""
Streams a chat request from the model.
@@ -160,12 +166,14 @@ class RunnerSupervisor:
raise RuntimeError("Runner process was found to be dead")
task_params = task.task_params
- assert isinstance(task_params, ChatCompletionTaskParams) # this is messy for now.
+ assert isinstance(
+ task_params, ChatCompletionTaskParams
+ ) # this is messy for now.
await self.write_queue.put(
ChatTaskMessage(
task_data=task_params,
),
- )
+ )
# This is simpler for now: we say 'request started' as soon as we've told runner to start, without waiting for an ack.
# If we need more reliability, the runner can have a new 'ready' message type.
@@ -175,13 +183,15 @@ class RunnerSupervisor:
prefil_timeout = get_prefil_timeout(self.model_shard_meta)
token_timeout = get_token_generate_timeout(self.model_shard_meta)
timeout = prefil_timeout
- self.logger.info(f'starting chat completion with timeout {timeout}')
+ self.logger.info(f"starting chat completion with timeout {timeout}")
while True:
try:
response = await self._read_with_error_check(timeout)
except asyncio.TimeoutError as e:
- self.logger.info(f'timed out from timeout duration {timeout} - {"prefil" if timeout == prefil_timeout else "decoding stage"}')
+ self.logger.info(
+ f"timed out from timeout duration {timeout} - {'prefil' if timeout == prefil_timeout else 'decoding stage'}"
+ )
raise e
match response:
@@ -199,17 +209,16 @@ class RunnerSupervisor:
break
case ErrorResponse():
await self.astop()
- raise RunnerError(response.error_type, response.error_message, response.traceback)
+ raise RunnerError(
+ response.error_type, response.error_message, response.traceback
+ )
case _:
- raise ValueError(f'Unexpected response type found: {response}')
+ raise ValueError(f"Unexpected response type found: {response}")
async def _write_coro(self):
while True:
message = await self.write_queue.get()
- await supervisor_write_message(
- self.runner_process,
- message
- )
+ await supervisor_write_message(self.runner_process, message)
async def _read_coro(self):
while True:
@@ -233,7 +242,6 @@ class RunnerSupervisor:
case _:
await self.read_queue.put(response)
-
async def astop(self) -> None:
# Cancel the stderr monitoring task
async def await_task(task: asyncio.Task[Any]):
@@ -241,14 +249,14 @@ class RunnerSupervisor:
task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await task
-
+
await await_task(self.stderr_task)
await await_task(self.read_task)
await await_task(self.write_task)
# Kill the process and all its children
await kill_process_tree(self.runner_process, self.logger)
-
+
# Wait to make sure that the model has been unloaded from memory
async def wait_for_memory_release() -> None:
required_memory_bytes = get_weights_size_kb(self.model_shard_meta) * 1024
@@ -258,7 +266,9 @@ class RunnerSupervisor:
if available_memory_bytes >= required_memory_bytes:
break
if asyncio.get_event_loop().time() - start_time > 30.0:
- self.logger.warning("Timeout waiting for memory release after 30 seconds")
+ self.logger.warning(
+ "Timeout waiting for memory release after 30 seconds"
+ )
break
await asyncio.sleep(0.1)
@@ -276,7 +286,9 @@ class RunnerSupervisor:
parent = psutil.Process(pid)
children = parent.children(recursive=True)
for child in reversed(children):
- with contextlib.suppress(psutil.NoSuchProcess, psutil.AccessDenied):
+ with contextlib.suppress(
+ psutil.NoSuchProcess, psutil.AccessDenied
+ ):
child.kill()
with contextlib.suppress(psutil.NoSuchProcess, psutil.AccessDenied):
parent.kill()
@@ -301,7 +313,7 @@ class RunnerSupervisor:
await self.astop()
# Accumulate all stderr messages from the queue
- stderr_output = ''
+ stderr_output = ""
while not self.stderr_queue.empty():
try:
line = self.stderr_queue.get_nowait()
@@ -312,7 +324,7 @@ class RunnerSupervisor:
# print('STDERR OUTPUT IS')
# print(stderr_output)
- self.logger.error(f'Error {self.runner_process.returncode}: {stderr_output}')
+ self.logger.error(f"Error {self.runner_process.returncode}: {stderr_output}")
return RunnerError(
error_type="MLXCrash",
error_message=stderr_output,
@@ -326,10 +338,10 @@ class RunnerSupervisor:
line_bytes = await self.runner_process.stderr.readline()
if not line_bytes:
break
- line = line_bytes.decode('utf-8').strip()
+ line = line_bytes.decode("utf-8").strip()
await self.stderr_queue.put(line)
self.logger.warning(f"Runner stderr read: {line}")
except Exception as e:
self.logger.warning(f"Error reading runner stderr: {e}")
- break
\ No newline at end of file
+ break
diff --git a/src/exo/worker/runner/utils.py b/src/exo/worker/runner/utils.py
index 2b04f424..c5c480ca 100644
--- a/src/exo/worker/runner/utils.py
+++ b/src/exo/worker/runner/utils.py
@@ -9,48 +9,57 @@ from exo.shared.constants import LB_DISK_GBPS, LB_MEMBW_GBPS, LB_TFLOPS
from exo.shared.types.worker.shards import ShardMetadata
-async def kill_process_tree(runner_process: asyncio.subprocess.Process, logger: Logger) -> None:
+async def kill_process_tree(
+ runner_process: asyncio.subprocess.Process, logger: Logger
+) -> None:
"""Kill the process and all its children forcefully."""
if runner_process.returncode is not None:
return # Process already dead
-
+
try:
# Get the main process
pid = runner_process.pid
-
+
# Find all child processes
try:
parent = psutil.Process(pid)
children = parent.children(recursive=True)
-
+
# Kill all children first (bottom-up)
for child in reversed(children):
with contextlib.suppress(psutil.NoSuchProcess, psutil.AccessDenied):
child.kill() # SIGKILL
-
+
# Kill the parent
with contextlib.suppress(psutil.NoSuchProcess, psutil.AccessDenied):
parent.kill() # SIGKILL
-
+
except psutil.NoSuchProcess:
# Process already gone, try subprocess kill anyway
runner_process.kill()
-
+
# Wait for the subprocess to exit
try:
await asyncio.wait_for(runner_process.wait(), timeout=2.0)
except asyncio.TimeoutError:
logger.error(f"Process {pid} did not exit after kill signal")
-
+
except Exception as e:
logger.error(f"Error killing process tree: {e}")
+
def get_runner_command() -> list[str]:
python = sys.executable
return [python, "-m", "exo.worker.runner.runner"]
+
def get_weights_size_kb(model_shard_meta: ShardMetadata) -> float:
- return (model_shard_meta.end_layer - model_shard_meta.start_layer) / model_shard_meta.n_layers * model_shard_meta.model_meta.storage_size_kilobytes
+ return (
+ (model_shard_meta.end_layer - model_shard_meta.start_layer)
+ / model_shard_meta.n_layers
+ * model_shard_meta.model_meta.storage_size_kilobytes
+ )
+
def get_init_timeout(model_shard_meta: ShardMetadata) -> float:
weights_size_kb = get_weights_size_kb(model_shard_meta)
@@ -59,17 +68,19 @@ def get_init_timeout(model_shard_meta: ShardMetadata) -> float:
return weights_size_kb / kbps_read + 2.0
+
def get_prefil_timeout(model_shard_meta: ShardMetadata) -> float:
weights_size_gb = get_weights_size_kb(model_shard_meta) / (1024 * 1024)
-
- tokens = 1000 # constant for now - the prompt is only tokenized in the device...
+
+ tokens = 1000 # constant for now - the prompt is only tokenized in the device...
prompt_gflops = tokens * weights_size_gb * 2
return LB_TFLOPS / (1024 * prompt_gflops) * 3 + 10.0
+
def get_token_generate_timeout(model_shard_meta: ShardMetadata) -> float:
weights_size_kb = get_weights_size_kb(model_shard_meta)
kbps_read = 1024 * 1024 * LB_MEMBW_GBPS / 3
- return weights_size_kb / kbps_read + 2.0
\ No newline at end of file
+ return weights_size_kb / kbps_read + 2.0
diff --git a/src/exo/worker/tests/__init__.py b/src/exo/worker/tests/__init__.py
index 0519ecba..e69de29b 100644
--- a/src/exo/worker/tests/__init__.py
+++ b/src/exo/worker/tests/__init__.py
@@ -1 +0,0 @@
-
\ No newline at end of file
diff --git a/src/exo/worker/tests/conftest.py b/src/exo/worker/tests/conftest.py
index e13624e3..3f24ae5c 100644
--- a/src/exo/worker/tests/conftest.py
+++ b/src/exo/worker/tests/conftest.py
@@ -33,25 +33,29 @@ def user_message():
"""Override this fixture in tests to customize the message"""
return "Hello, how are you?"
+
@pytest.fixture
def logger() -> Logger:
import logging
+
logger = getLogger("test_logger")
logger.setLevel(logging.DEBUG)
-
+
# Add console handler if none exists
if not logger.handlers:
handler = logging.StreamHandler()
handler.setLevel(logging.DEBUG)
- formatter = logging.Formatter('%(name)s - %(levelname)s - %(message)s')
+ formatter = logging.Formatter("%(name)s - %(levelname)s - %(message)s")
handler.setFormatter(formatter)
logger.addHandler(handler)
-
+
return logger
+
@pytest.fixture
async def model_meta() -> ModelMetadata:
- return await get_model_meta('mlx-community/Llama-3.2-1B-Instruct-4bit')
+ return await get_model_meta("mlx-community/Llama-3.2-1B-Instruct-4bit")
+
@pytest.fixture
def hosts():
@@ -66,8 +70,11 @@ def hosts():
return _hosts
+
@pytest.fixture
-def pipeline_shard_meta(model_meta: ModelMetadata) -> Callable[[int, int], PipelineShardMetadata]:
+def pipeline_shard_meta(
+ model_meta: ModelMetadata,
+) -> Callable[[int, int], PipelineShardMetadata]:
def _pipeline_shard_meta(
num_nodes: int = 1, device_rank: int = 0
) -> PipelineShardMetadata:
@@ -91,8 +98,12 @@ def pipeline_shard_meta(model_meta: ModelMetadata) -> Callable[[int, int], Pipel
return _pipeline_shard_meta
+
@pytest.fixture
-def instance(pipeline_shard_meta: Callable[[int, int], PipelineShardMetadata], hosts: Callable[[int], list[Host]]):
+def instance(
+ pipeline_shard_meta: Callable[[int, int], PipelineShardMetadata],
+ hosts: Callable[[int], list[Host]],
+):
from typing import Optional
def _instance(
@@ -108,20 +119,20 @@ def instance(pipeline_shard_meta: Callable[[int, int], PipelineShardMetadata], h
shard_assignments = ShardAssignments(
model_id=resolved_model_id,
- runner_to_shard={
- resolved_runner_id: pipeline_shard_meta(1, 0)
- },
- node_to_runner={resolved_node_id: resolved_runner_id}
+ runner_to_shard={resolved_runner_id: pipeline_shard_meta(1, 0)},
+ node_to_runner={resolved_node_id: resolved_runner_id},
)
return Instance(
instance_id=resolved_instance_id,
instance_type=InstanceStatus.ACTIVE,
shard_assignments=shard_assignments,
- hosts=hosts(1)
+ hosts=hosts(1),
)
+
return _instance
+
@pytest.fixture
def completion_create_params(user_message: str) -> ChatCompletionTaskParams:
return ChatCompletionTaskParams(
@@ -130,12 +141,13 @@ def completion_create_params(user_message: str) -> ChatCompletionTaskParams:
stream=True,
)
+
@pytest.fixture
def chat_completion_task(completion_create_params: ChatCompletionTaskParams):
def _chat_completion_task(
instance_id: Optional[InstanceId] = None,
task_id: Optional[TaskId] = None,
- user_message: str = "Hello"
+ user_message: str = "Hello",
) -> ChatCompletionTask:
resolved_instance_id = instance_id if instance_id is not None else INSTANCE_1_ID
resolved_task_id = task_id if task_id is not None else TASK_1_ID
@@ -145,6 +157,7 @@ def chat_completion_task(completion_create_params: ChatCompletionTaskParams):
instance_id=resolved_instance_id,
task_type=TaskType.CHAT_COMPLETION,
task_status=TaskStatus.PENDING,
- task_params=completion_create_params
+ task_params=completion_create_params,
)
+
return _chat_completion_task
diff --git a/src/exo/worker/tests/constants.py b/src/exo/worker/tests/constants.py
index 49ff6876..4de842f5 100644
--- a/src/exo/worker/tests/constants.py
+++ b/src/exo/worker/tests/constants.py
@@ -16,11 +16,11 @@ RUNNER_2_ID: Final[RunnerId] = RunnerId("33333333-3333-4333-8333-333333333333")
INSTANCE_1_ID: Final[InstanceId] = InstanceId("22222222-2222-4222-8222-222222222222")
INSTANCE_2_ID: Final[InstanceId] = InstanceId("44444444-4444-4444-8444-444444444444")
-MODEL_A_ID: Final[ModelId] = 'mlx-community/Llama-3.2-1B-Instruct-4bit'
-MODEL_B_ID: Final[ModelId] = 'mlx-community/TinyLlama-1.1B-Chat-v1.0'
+MODEL_A_ID: Final[ModelId] = "mlx-community/Llama-3.2-1B-Instruct-4bit"
+MODEL_B_ID: Final[ModelId] = "mlx-community/TinyLlama-1.1B-Chat-v1.0"
TASK_1_ID: Final[TaskId] = TaskId("55555555-5555-4555-8555-555555555555")
TASK_2_ID: Final[TaskId] = TaskId("66666666-6666-4666-8666-666666666666")
COMMAND_1_ID: Final[CommandId] = CommandId("77777777-7777-4777-8777-777777777777")
-COMMAND_2_ID: Final[CommandId] = CommandId("88888888-8888-4888-8888-888888888888")
\ No newline at end of file
+COMMAND_2_ID: Final[CommandId] = CommandId("88888888-8888-4888-8888-888888888888")
diff --git a/src/exo/worker/tests/test_download.py b/src/exo/worker/tests/test_download.py
index 6331562b..3ce6b964 100644
--- a/src/exo/worker/tests/test_download.py
+++ b/src/exo/worker/tests/test_download.py
@@ -10,7 +10,9 @@ from exo.worker.download.shard_downloader import ShardDownloader
@pytest.mark.slow
@pytest.mark.asyncio
-async def test_shard_downloader(pipeline_shard_meta: Callable[[int, int], PipelineShardMetadata]):
+async def test_shard_downloader(
+ pipeline_shard_meta: Callable[[int, int], PipelineShardMetadata],
+):
shard_downloader: ShardDownloader = exo_shard_downloader()
shard_downloader.on_progress(
lambda shard, progress: print(f"Download progress: {progress}")
diff --git a/src/exo/worker/tests/test_handlers/conftest.py b/src/exo/worker/tests/test_handlers/conftest.py
index a6d96ef6..7707754a 100644
--- a/src/exo/worker/tests/test_handlers/conftest.py
+++ b/src/exo/worker/tests/test_handlers/conftest.py
@@ -28,17 +28,26 @@ async def worker(logger: Logger):
shard_downloader = NoopShardDownloader()
await event_log_manager.initialize()
- return Worker(NODE_A, logger, shard_downloader, worker_events=event_log_manager.global_events, global_events=event_log_manager.global_events)
+ return Worker(
+ NODE_A,
+ logger,
+ shard_downloader,
+ worker_events=event_log_manager.global_events,
+ global_events=event_log_manager.global_events,
+ )
+
# TODO: instance_id and runner_id are selectable.
@pytest.fixture
-async def worker_with_assigned_runner(worker: Worker, instance: Callable[[InstanceId, NodeId, RunnerId], Instance]):
+async def worker_with_assigned_runner(
+ worker: Worker, instance: Callable[[InstanceId, NodeId, RunnerId], Instance]
+):
"""Fixture that provides a worker with an already assigned runner."""
-
+
instance_id = INSTANCE_1_ID
runner_id = RUNNER_1_ID
instance_obj: Instance = instance(instance_id, worker.node_id, runner_id)
-
+
# Assign the runner
assign_op = AssignRunnerOp(
runner_id=runner_id,
@@ -46,14 +55,17 @@ async def worker_with_assigned_runner(worker: Worker, instance: Callable[[Instan
hosts=instance_obj.hosts,
instance_id=instance_obj.instance_id,
)
-
+
async for _ in worker.execute_op(assign_op):
pass
-
+
return worker, instance_obj
+
@pytest.fixture
-async def worker_with_running_runner(worker_with_assigned_runner: tuple[Worker, Instance]):
+async def worker_with_running_runner(
+ worker_with_assigned_runner: tuple[Worker, Instance],
+):
"""Fixture that provides a worker with an already assigned runner."""
worker, instance_obj = worker_with_assigned_runner
@@ -67,4 +79,3 @@ async def worker_with_running_runner(worker_with_assigned_runner: tuple[Worker,
assert supervisor.healthy
return worker, instance_obj
-
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 a11750a5..a58ecd37 100644
--- a/src/exo/worker/tests/test_handlers/test_handlers_happy.py
+++ b/src/exo/worker/tests/test_handlers/test_handlers_happy.py
@@ -34,7 +34,9 @@ from exo.worker.tests.test_handlers.utils import read_events_op
@pytest.mark.asyncio
-async def test_assign_op(worker: Worker, instance: Callable[[InstanceId, NodeId, RunnerId], Instance]):
+async def test_assign_op(
+ worker: Worker, instance: Callable[[InstanceId, NodeId, RunnerId], Instance]
+):
instance_obj: Instance = instance(InstanceId(), worker.node_id, RUNNER_1_ID)
assign_op = AssignRunnerOp(
@@ -57,13 +59,12 @@ async def test_assign_op(worker: Worker, instance: Callable[[InstanceId, NodeId,
assert RUNNER_1_ID in worker.assigned_runners
assert isinstance(worker.assigned_runners[RUNNER_1_ID].status, InactiveRunnerStatus)
+
@pytest.mark.asyncio
async def test_unassign_op(worker_with_assigned_runner: tuple[Worker, Instance]):
worker, _ = worker_with_assigned_runner
- unassign_op = UnassignRunnerOp(
- runner_id=RUNNER_1_ID
- )
+ unassign_op = UnassignRunnerOp(runner_id=RUNNER_1_ID)
events = await read_events_op(worker, unassign_op)
@@ -72,11 +73,12 @@ async def test_unassign_op(worker_with_assigned_runner: tuple[Worker, Instance])
assert len(events) == 1
assert isinstance(events[0], RunnerDeleted)
+
@pytest.mark.asyncio
async def test_runner_up_op(
- worker_with_assigned_runner: tuple[Worker, Instance],
- chat_completion_task: Callable[[], ChatCompletionTask],
- ):
+ worker_with_assigned_runner: tuple[Worker, Instance],
+ chat_completion_task: Callable[[], ChatCompletionTask],
+):
worker, _ = worker_with_assigned_runner
runner_up_op = RunnerUpOp(runner_id=RUNNER_1_ID)
@@ -92,7 +94,7 @@ async def test_runner_up_op(
assert supervisor is not None
assert supervisor.healthy
- full_response = ''
+ full_response = ""
async for chunk in supervisor.stream_response(task=chat_completion_task()):
if isinstance(chunk, TokenChunk):
@@ -104,7 +106,8 @@ async def test_runner_up_op(
runner = worker.assigned_runners[RUNNER_1_ID].runner
assert runner is not None
- await runner.astop() # Neat cleanup.
+ await runner.astop() # Neat cleanup.
+
@pytest.mark.asyncio
async def test_runner_down_op(worker_with_running_runner: tuple[Worker, Instance]):
@@ -117,43 +120,47 @@ async def test_runner_down_op(worker_with_running_runner: tuple[Worker, Instance
assert isinstance(events[0], RunnerStatusUpdated)
assert isinstance(events[0].runner_status, InactiveRunnerStatus)
+
@pytest.mark.asyncio
async def test_execute_task_op(
- worker_with_running_runner: tuple[Worker, Instance],
- chat_completion_task: Callable[[], ChatCompletionTask]):
+ worker_with_running_runner: tuple[Worker, Instance],
+ chat_completion_task: Callable[[], ChatCompletionTask],
+):
worker, _ = worker_with_running_runner
- execute_task_op = ExecuteTaskOp(
- runner_id=RUNNER_1_ID,
- task=chat_completion_task()
- )
+ execute_task_op = ExecuteTaskOp(runner_id=RUNNER_1_ID, task=chat_completion_task())
events = await read_events_op(worker, execute_task_op)
assert len(events) > 20
- print(f'{events=}')
-
+ print(f"{events=}")
assert isinstance(events[0], RunnerStatusUpdated)
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(events[-1].runner_status, LoadedRunnerStatus) # It should not have failed.
-
- gen_events: list[ChunkGenerated] = [x for x in events if isinstance(x, ChunkGenerated)]
- text_chunks: list[TokenChunk] = [x.chunk for x in gen_events if isinstance(x.chunk, TokenChunk)]
+ assert isinstance(
+ events[-1].runner_status, LoadedRunnerStatus
+ ) # It should not have failed.
+
+ gen_events: list[ChunkGenerated] = [
+ x for x in events if isinstance(x, ChunkGenerated)
+ ]
+ text_chunks: list[TokenChunk] = [
+ x.chunk for x in gen_events if isinstance(x.chunk, TokenChunk)
+ ]
assert len(text_chunks) == len(events) - 4
-
- output_text = ''.join([x.text for x in text_chunks])
- assert '42' in output_text
+
+ output_text = "".join([x.text for x in text_chunks])
+ assert "42" in output_text
runner = worker.assigned_runners[RUNNER_1_ID].runner
assert runner is not None
- await runner.astop() # Neat cleanup.
+ await runner.astop() # Neat cleanup.
diff --git a/src/exo/worker/tests/test_handlers/test_handlers_sad.py b/src/exo/worker/tests/test_handlers/test_handlers_sad.py
index 588cee8a..97d2772c 100644
--- a/src/exo/worker/tests/test_handlers/test_handlers_sad.py
+++ b/src/exo/worker/tests/test_handlers/test_handlers_sad.py
@@ -19,8 +19,9 @@ from exo.worker.tests.test_handlers.utils import read_events_op
@pytest.mark.asyncio
async def test_runner_up_fails(
- worker_with_assigned_runner: tuple[Worker, Instance],
- chat_completion_task: Callable[[], ChatCompletionTask]):
+ worker_with_assigned_runner: tuple[Worker, Instance],
+ chat_completion_task: Callable[[], ChatCompletionTask],
+):
worker, _ = worker_with_assigned_runner
worker.assigned_runners[RUNNER_1_ID].shard_metadata.immediate_exception = True
@@ -29,10 +30,12 @@ async def test_runner_up_fails(
with pytest.raises(RunnerError):
await read_events_op(worker, runner_up_op)
+
@pytest.mark.asyncio
async def test_runner_up_timeouts(
- worker_with_assigned_runner: tuple[Worker, Instance],
- chat_completion_task: Callable[[], ChatCompletionTask]):
+ worker_with_assigned_runner: tuple[Worker, Instance],
+ chat_completion_task: Callable[[], ChatCompletionTask],
+):
worker, _ = worker_with_assigned_runner
worker.assigned_runners[RUNNER_1_ID].shard_metadata.should_timeout = 10
@@ -41,42 +44,40 @@ async def test_runner_up_timeouts(
with pytest.raises(asyncio.TimeoutError):
await read_events_op(worker, runner_up_op)
+
@pytest.mark.asyncio
async def test_execute_task_fails(
- worker_with_running_runner: tuple[Worker, Instance],
- chat_completion_task: Callable[[], ChatCompletionTask]):
+ worker_with_running_runner: tuple[Worker, Instance],
+ chat_completion_task: Callable[[], ChatCompletionTask],
+):
worker, _ = worker_with_running_runner
task = chat_completion_task()
messages = task.task_params.messages
- messages[0].content = 'Artificial prompt: EXO RUNNER MUST FAIL'
+ messages[0].content = "Artificial prompt: EXO RUNNER MUST FAIL"
- execute_task_op = ExecuteTaskOp(
- runner_id=RUNNER_1_ID,
- task=task
- )
+ execute_task_op = ExecuteTaskOp(runner_id=RUNNER_1_ID, task=task)
with pytest.raises(RunnerError):
await read_events_op(worker, execute_task_op)
+
@pytest.mark.asyncio
async def test_execute_task_timeouts(
- worker_with_running_runner: tuple[Worker, Instance],
- chat_completion_task: Callable[[], ChatCompletionTask]):
+ worker_with_running_runner: tuple[Worker, Instance],
+ chat_completion_task: Callable[[], ChatCompletionTask],
+):
worker, _ = worker_with_running_runner
task = chat_completion_task()
messages = task.task_params.messages
- messages[0].content = 'Artificial prompt: EXO RUNNER MUST TIMEOUT'
+ messages[0].content = "Artificial prompt: EXO RUNNER MUST TIMEOUT"
- execute_task_op = ExecuteTaskOp(
- runner_id=RUNNER_1_ID,
- task=task
- )
+ execute_task_op = ExecuteTaskOp(runner_id=RUNNER_1_ID, task=task)
with pytest.raises(asyncio.TimeoutError):
await read_events_op(worker, execute_task_op)
# TODO: Much more to do here!
-# runner assigned download stuff
\ No newline at end of file
+# runner assigned download stuff
diff --git a/src/exo/worker/tests/test_handlers/utils.py b/src/exo/worker/tests/test_handlers/utils.py
index 4b095342..db5af33a 100644
--- a/src/exo/worker/tests/test_handlers/utils.py
+++ b/src/exo/worker/tests/test_handlers/utils.py
@@ -1,7 +1,6 @@
## Tests for worker state handlers
-
from exo.shared.types.events import (
Event,
)
@@ -15,4 +14,4 @@ async def read_events_op(worker: Worker, op: RunnerOp) -> list[Event]:
events: list[Event] = []
async for event in worker.execute_op(op):
events.append(event)
- return events
\ No newline at end of file
+ return events
diff --git a/src/exo/worker/tests/test_integration/conftest.py b/src/exo/worker/tests/test_integration/conftest.py
index 4e00d414..2f1888ec 100644
--- a/src/exo/worker/tests/test_integration/conftest.py
+++ b/src/exo/worker/tests/test_integration/conftest.py
@@ -19,8 +19,12 @@ def user_message():
@pytest.fixture
-def worker_running(logger: Logger) -> Callable[[NodeId], Awaitable[tuple[Worker, AsyncSQLiteEventStorage]]]:
- async def _worker_running(node_id: NodeId) -> tuple[Worker, AsyncSQLiteEventStorage]:
+def worker_running(
+ logger: Logger,
+) -> Callable[[NodeId], Awaitable[tuple[Worker, AsyncSQLiteEventStorage]]]:
+ async def _worker_running(
+ node_id: NodeId,
+ ) -> tuple[Worker, AsyncSQLiteEventStorage]:
event_log_manager = EventLogManager(EventLogConfig(), logger)
await event_log_manager.initialize()
@@ -28,9 +32,15 @@ def worker_running(logger: Logger) -> Callable[[NodeId], Awaitable[tuple[Worker,
await global_events.delete_all_events()
shard_downloader = NoopShardDownloader()
- worker = Worker(node_id, logger=logger, shard_downloader=shard_downloader, worker_events=global_events, global_events=global_events)
+ worker = Worker(
+ node_id,
+ logger=logger,
+ shard_downloader=shard_downloader,
+ worker_events=global_events,
+ global_events=global_events,
+ )
asyncio.create_task(run(worker, logger))
return worker, global_events
- return _worker_running
\ No newline at end of file
+ return _worker_running
diff --git a/src/exo/worker/tests/test_integration/integration_utils.py b/src/exo/worker/tests/test_integration/integration_utils.py
index c059613a..1112dbd2 100644
--- a/src/exo/worker/tests/test_integration/integration_utils.py
+++ b/src/exo/worker/tests/test_integration/integration_utils.py
@@ -1,5 +1,3 @@
-
-
import asyncio
from typing import Callable, Optional, Tuple, TypeVar
@@ -9,10 +7,12 @@ from exo.shared.types.events.chunks import TokenChunk
from exo.shared.types.tasks import TaskId, TaskStatus
-async def read_streaming_response(global_events: AsyncSQLiteEventStorage, filter_task: Optional[TaskId] = None) -> Tuple[bool, bool, str]:
+async def read_streaming_response(
+ global_events: AsyncSQLiteEventStorage, filter_task: Optional[TaskId] = None
+) -> Tuple[bool, bool, str]:
# Read off all events - these should be our GenerationChunk events
seen_task_started, seen_task_finished = 0, 0
- response_string = ''
+ response_string = ""
finish_reason: str | None = None
if not filter_task:
@@ -22,14 +22,18 @@ async def read_streaming_response(global_events: AsyncSQLiteEventStorage, filter
idx = 0
while not found:
events = await global_events.get_events_since(idx)
-
+
for event in events:
- if isinstance(event.event, TaskStateUpdated) and event.event.task_status == TaskStatus.RUNNING and event.event.task_id == filter_task:
+ if (
+ isinstance(event.event, TaskStateUpdated)
+ and event.event.task_status == TaskStatus.RUNNING
+ and event.event.task_id == filter_task
+ ):
found = True
idx = event.idx_in_log - 1
break
- print(f'START IDX {idx}')
+ print(f"START IDX {idx}")
while not finish_reason:
events = await global_events.get_events_since(idx)
@@ -54,12 +58,14 @@ async def read_streaming_response(global_events: AsyncSQLiteEventStorage, filter
await asyncio.sleep(0.2)
- print(f'event log: {await global_events.get_events_since(0)}')
+ print(f"event log: {await global_events.get_events_since(0)}")
return seen_task_started == 1, seen_task_finished == 1, response_string
+
T = TypeVar("T")
+
async def until_event_with_timeout(
global_events: AsyncSQLiteEventStorage,
event_type: type[T],
@@ -72,10 +78,12 @@ async def until_event_with_timeout(
events = await global_events.get_events_since(idx)
if events:
for wrapped_event in events:
- if isinstance(wrapped_event.event, event_type) and condition(wrapped_event.event):
+ if isinstance(wrapped_event.event, event_type) and condition(
+ wrapped_event.event
+ ):
times_seen += 1
if times_seen >= multiplicity:
return
idx = events[-1].idx_in_log
- await asyncio.sleep(0.01)
\ No newline at end of file
+ await asyncio.sleep(0.01)
diff --git a/src/exo/worker/tests/test_integration/test_inference.py b/src/exo/worker/tests/test_integration/test_inference.py
index bb2c5966..8262af4a 100644
--- a/src/exo/worker/tests/test_integration/test_inference.py
+++ b/src/exo/worker/tests/test_integration/test_inference.py
@@ -13,7 +13,13 @@ from exo.shared.types.events import (
TaskCreated,
)
from exo.shared.types.models import ModelId
-from exo.shared.types.tasks import ChatCompletionTask, Task, TaskId, TaskStatus, TaskType
+from exo.shared.types.tasks import (
+ ChatCompletionTask,
+ Task,
+ TaskId,
+ TaskStatus,
+ TaskType,
+)
from exo.shared.types.worker.common import InstanceId, RunnerId
from exo.shared.types.worker.instances import (
Instance,
@@ -39,10 +45,12 @@ from exo.worker.worker import Worker
async def test_runner_inference(
- worker_running: Callable[[NodeId], Awaitable[tuple[Worker, AsyncSQLiteEventStorage]]],
+ worker_running: Callable[
+ [NodeId], Awaitable[tuple[Worker, AsyncSQLiteEventStorage]]
+ ],
instance: Callable[[InstanceId, NodeId, RunnerId], Instance],
- chat_completion_task: Callable[[InstanceId, TaskId], Task]
- ):
+ chat_completion_task: Callable[[InstanceId, TaskId], Task],
+):
_worker, global_events = await worker_running(NODE_A)
instance_value: Instance = instance(INSTANCE_1_ID, NODE_A, RUNNER_1_ID)
@@ -54,38 +62,40 @@ async def test_runner_inference(
InstanceCreated(
instance=instance_value,
),
- TaskCreated(
- task_id=task.task_id,
- task=task
- )
- ],
- origin=MASTER_NODE_ID
+ TaskCreated(task_id=task.task_id, task=task),
+ ],
+ origin=MASTER_NODE_ID,
)
-
+
# TODO: This needs to get fixed - sometimes it misses the 'starting' event.
- seen_task_started, seen_task_finished, response_string = await read_streaming_response(global_events)
+ (
+ seen_task_started,
+ seen_task_finished,
+ response_string,
+ ) = await read_streaming_response(global_events)
assert seen_task_started
assert seen_task_finished
- assert 'tokyo' in response_string.lower()
+ assert "tokyo" in response_string.lower()
await global_events.append_events(
[
InstanceDeleted(
instance_id=instance_value.instance_id,
),
- ],
- origin=MASTER_NODE_ID
+ ],
+ origin=MASTER_NODE_ID,
)
await asyncio.sleep(0.3)
+
async def test_2_runner_inference(
logger: Logger,
pipeline_shard_meta: Callable[[int, int], PipelineShardMetadata],
hosts: Callable[[int], list[Host]],
- chat_completion_task: Callable[[InstanceId, TaskId], Task]
- ):
+ chat_completion_task: Callable[[InstanceId, TaskId], Task],
+):
event_log_manager = EventLogManager(EventLogConfig(), logger)
await event_log_manager.initialize()
shard_downloader = NoopShardDownloader()
@@ -93,54 +103,61 @@ async def test_2_runner_inference(
global_events = event_log_manager.global_events
await global_events.delete_all_events()
- worker1 = Worker(NODE_A, logger=logger, shard_downloader=shard_downloader, worker_events=global_events, global_events=global_events)
+ worker1 = Worker(
+ NODE_A,
+ logger=logger,
+ shard_downloader=shard_downloader,
+ worker_events=global_events,
+ global_events=global_events,
+ )
asyncio.create_task(run(worker1, logger))
- worker2 = Worker(NODE_B, logger=logger, shard_downloader=shard_downloader, worker_events=global_events, global_events=global_events)
+ worker2 = Worker(
+ NODE_B,
+ logger=logger,
+ shard_downloader=shard_downloader,
+ worker_events=global_events,
+ global_events=global_events,
+ )
asyncio.create_task(run(worker2, logger))
## Instance
- model_id = ModelId('mlx-community/Llama-3.2-1B-Instruct-4bit')
+ model_id = ModelId("mlx-community/Llama-3.2-1B-Instruct-4bit")
shard_assignments = ShardAssignments(
model_id=model_id,
runner_to_shard={
RUNNER_1_ID: pipeline_shard_meta(2, 0),
- RUNNER_2_ID: pipeline_shard_meta(2, 1)
+ RUNNER_2_ID: pipeline_shard_meta(2, 1),
},
- node_to_runner={
- NODE_A: RUNNER_1_ID,
- NODE_B: RUNNER_2_ID
- }
+ node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID},
)
-
+
instance = Instance(
instance_id=INSTANCE_1_ID,
instance_type=InstanceStatus.ACTIVE,
shard_assignments=shard_assignments,
- hosts=hosts(2)
+ hosts=hosts(2),
)
task = chat_completion_task(INSTANCE_1_ID, TASK_1_ID)
await global_events.append_events(
[
- InstanceCreated(
- instance=instance
- ),
- TaskCreated(
- task_id=task.task_id,
- task=task
- )
- ],
- origin=MASTER_NODE_ID
+ InstanceCreated(instance=instance),
+ TaskCreated(task_id=task.task_id, task=task),
+ ],
+ origin=MASTER_NODE_ID,
)
- seen_task_started, seen_task_finished, response_string = await read_streaming_response(global_events)
+ (
+ seen_task_started,
+ seen_task_finished,
+ response_string,
+ ) = await read_streaming_response(global_events)
assert seen_task_started
assert seen_task_finished
- assert 'tokyo' in response_string.lower()
-
+ assert "tokyo" in response_string.lower()
idx = await global_events.get_last_idx()
await asyncio.sleep(1.0)
@@ -152,18 +169,19 @@ async def test_2_runner_inference(
InstanceDeleted(
instance_id=instance.instance_id,
),
- ],
- origin=MASTER_NODE_ID
+ ],
+ origin=MASTER_NODE_ID,
)
await asyncio.sleep(2.0)
+
# TODO: Multi message parallel
async def test_2_runner_multi_message(
logger: Logger,
pipeline_shard_meta: Callable[[int, int], PipelineShardMetadata],
hosts: Callable[[int], list[Host]],
- ):
+):
event_log_manager = EventLogManager(EventLogConfig(), logger)
await event_log_manager.initialize()
shard_downloader = NoopShardDownloader()
@@ -171,32 +189,41 @@ async def test_2_runner_multi_message(
global_events = event_log_manager.global_events
await global_events.delete_all_events()
- worker1 = Worker(NODE_A, logger=logger, shard_downloader=shard_downloader, worker_events=global_events, global_events=global_events)
+ worker1 = Worker(
+ NODE_A,
+ logger=logger,
+ shard_downloader=shard_downloader,
+ worker_events=global_events,
+ global_events=global_events,
+ )
asyncio.create_task(run(worker1, logger))
- worker2 = Worker(NODE_B, logger=logger, shard_downloader=shard_downloader, worker_events=global_events, global_events=global_events)
+ worker2 = Worker(
+ NODE_B,
+ logger=logger,
+ shard_downloader=shard_downloader,
+ worker_events=global_events,
+ global_events=global_events,
+ )
asyncio.create_task(run(worker2, logger))
## Instance
- model_id = ModelId('mlx-community/Llama-3.2-1B-Instruct-4bit')
+ model_id = ModelId("mlx-community/Llama-3.2-1B-Instruct-4bit")
shard_assignments = ShardAssignments(
model_id=model_id,
runner_to_shard={
RUNNER_1_ID: pipeline_shard_meta(2, 0),
- RUNNER_2_ID: pipeline_shard_meta(2, 1)
+ RUNNER_2_ID: pipeline_shard_meta(2, 1),
},
- node_to_runner={
- NODE_A: RUNNER_1_ID,
- NODE_B: RUNNER_2_ID
- }
+ node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID},
)
-
+
instance = Instance(
instance_id=INSTANCE_1_ID,
instance_type=InstanceStatus.ACTIVE,
shard_assignments=shard_assignments,
- hosts=hosts(2)
+ hosts=hosts(2),
)
# Task - we have three messages here, which is what the task is about
@@ -204,9 +231,16 @@ async def test_2_runner_multi_message(
completion_create_params = ChatCompletionTaskParams(
model="gpt-4",
messages=[
- ChatCompletionMessage(role="user", content='What is the capital of France?'),
- ChatCompletionMessage(role="assistant", content='The capital of France is Paris.'),
- ChatCompletionMessage(role="user", content='Ok great. Now write me a haiku about what you can do there.'),
+ ChatCompletionMessage(
+ role="user", content="What is the capital of France?"
+ ),
+ ChatCompletionMessage(
+ role="assistant", content="The capital of France is Paris."
+ ),
+ ChatCompletionMessage(
+ role="user",
+ content="Ok great. Now write me a haiku about what you can do there.",
+ ),
],
stream=True,
)
@@ -217,28 +251,29 @@ async def test_2_runner_multi_message(
instance_id=INSTANCE_1_ID,
task_type=TaskType.CHAT_COMPLETION,
task_status=TaskStatus.PENDING,
- task_params=completion_create_params
+ task_params=completion_create_params,
)
await global_events.append_events(
[
- InstanceCreated(
- instance=instance
- ),
- TaskCreated(
- task_id=task.task_id,
- task=task
- )
- ],
- origin=MASTER_NODE_ID
+ InstanceCreated(instance=instance),
+ TaskCreated(task_id=task.task_id, task=task),
+ ],
+ origin=MASTER_NODE_ID,
)
- seen_task_started, seen_task_finished, response_string = await read_streaming_response(global_events)
+ (
+ seen_task_started,
+ seen_task_finished,
+ response_string,
+ ) = await read_streaming_response(global_events)
assert seen_task_started
assert seen_task_finished
- assert any(keyword in response_string.lower() for keyword in ('kiss', 'paris', 'art', 'love'))
-
+ assert any(
+ keyword in response_string.lower()
+ for keyword in ("kiss", "paris", "art", "love")
+ )
idx = await global_events.get_last_idx()
await asyncio.sleep(1.0)
@@ -250,8 +285,8 @@ async def test_2_runner_multi_message(
InstanceDeleted(
instance_id=instance.instance_id,
),
- ],
- origin=MASTER_NODE_ID
+ ],
+ origin=MASTER_NODE_ID,
)
await asyncio.sleep(2.0)
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 8e2d25fa..d5aa4688 100644
--- a/src/exo/worker/tests/test_integration/test_inference_sad.py
+++ b/src/exo/worker/tests/test_integration/test_inference_sad.py
@@ -46,42 +46,64 @@ def user_message():
async def test_stream_response_failed_always(
monkeypatch: MonkeyPatch,
- worker_running: Callable[[NodeId], Awaitable[tuple[Worker, AsyncSQLiteEventStorage]]],
+ worker_running: Callable[
+ [NodeId], Awaitable[tuple[Worker, AsyncSQLiteEventStorage]]
+ ],
instance: Callable[[InstanceId, NodeId, RunnerId], Instance],
- chat_completion_task: Callable[[InstanceId, TaskId], Task]
+ chat_completion_task: Callable[[InstanceId, TaskId], Task],
) -> None:
_, global_events = await worker_running(NODE_A)
instance_value: Instance = instance(INSTANCE_1_ID, NODE_A, RUNNER_1_ID)
instance_value.instance_type = InstanceStatus.ACTIVE
-
+
async def mock_stream_response(
self: RunnerSupervisor,
task: Task,
- request_started_callback: Callable[..., CoroutineType[Any, Any, None]] | None = None,
+ request_started_callback: Callable[..., CoroutineType[Any, Any, None]]
+ | None = None,
) -> AsyncGenerator[GenerationChunk]:
raise RuntimeError("Simulated stream response failure")
return
- yield
-
- monkeypatch.setattr(RunnerSupervisor, 'stream_response', mock_stream_response)
+ yield
+
+ monkeypatch.setattr(RunnerSupervisor, "stream_response", mock_stream_response)
task: Task = chat_completion_task(INSTANCE_1_ID, TASK_1_ID)
await global_events.append_events(
[
InstanceCreated(instance=instance_value),
- TaskCreated(task_id=task.task_id, task=task)
- ],
- origin=MASTER_NODE_ID
+ TaskCreated(task_id=task.task_id, task=task),
+ ],
+ origin=MASTER_NODE_ID,
)
await until_event_with_timeout(global_events, InstanceDeleted)
-
events = await global_events.get_events_since(0)
- assert len([x for x in events if isinstance(x.event, RunnerStatusUpdated) and isinstance(x.event.runner_status, FailedRunnerStatus)]) == 3
- assert len([x for x in events if isinstance(x.event, TaskStateUpdated) and x.event.task_status == TaskStatus.FAILED]) == 3
+ assert (
+ len(
+ [
+ x
+ for x in events
+ if isinstance(x.event, RunnerStatusUpdated)
+ and isinstance(x.event.runner_status, FailedRunnerStatus)
+ ]
+ )
+ == 3
+ )
+ assert (
+ len(
+ [
+ x
+ for x in events
+ if isinstance(x.event, TaskStateUpdated)
+ and x.event.task_status == TaskStatus.FAILED
+ ]
+ )
+ == 3
+ )
assert any([isinstance(x.event, InstanceDeleted) for x in events])
await global_events.append_events(
@@ -89,17 +111,20 @@ async def test_stream_response_failed_always(
InstanceDeleted(
instance_id=instance_value.instance_id,
),
- ],
- origin=MASTER_NODE_ID
+ ],
+ origin=MASTER_NODE_ID,
)
await asyncio.sleep(0.3)
+
async def test_stream_response_failed_once(
monkeypatch: MonkeyPatch,
- worker_running: Callable[[NodeId], Awaitable[tuple[Worker, AsyncSQLiteEventStorage]]],
+ worker_running: Callable[
+ [NodeId], Awaitable[tuple[Worker, AsyncSQLiteEventStorage]]
+ ],
instance: Callable[[InstanceId, NodeId, RunnerId], Instance],
- chat_completion_task: Callable[[InstanceId, TaskId], Task]
+ chat_completion_task: Callable[[InstanceId, TaskId], Task],
):
failed_already = False
original_stream_response = RunnerSupervisor.stream_response
@@ -107,19 +132,22 @@ async def test_stream_response_failed_once(
async def mock_stream_response(
self: RunnerSupervisor,
task: Task,
- request_started_callback: Callable[..., CoroutineType[Any, Any, None]] | None = None,
+ request_started_callback: Callable[..., CoroutineType[Any, Any, None]]
+ | None = None,
) -> AsyncGenerator[GenerationChunk]:
nonlocal failed_already
if not failed_already:
failed_already = True
raise RuntimeError("Simulated stream response failure")
else:
- async for event in original_stream_response(self, task, request_started_callback):
+ async for event in original_stream_response(
+ self, task, request_started_callback
+ ):
yield event
return
-
- monkeypatch.setattr(RunnerSupervisor, 'stream_response', mock_stream_response)
-
+
+ monkeypatch.setattr(RunnerSupervisor, "stream_response", mock_stream_response)
+
worker, global_events = await worker_running(NODE_A)
instance_value: Instance = instance(INSTANCE_1_ID, NODE_A, RUNNER_1_ID)
@@ -129,26 +157,52 @@ async def test_stream_response_failed_once(
await global_events.append_events(
[
InstanceCreated(instance=instance_value),
- TaskCreated(task_id=task.task_id, task=task)
- ],
- origin=MASTER_NODE_ID
+ TaskCreated(task_id=task.task_id, task=task),
+ ],
+ origin=MASTER_NODE_ID,
+ )
+
+ await until_event_with_timeout(
+ global_events,
+ ChunkGenerated,
+ 1,
+ condition=lambda x: isinstance(x.chunk, TokenChunk)
+ and x.chunk.finish_reason is not None,
)
- await until_event_with_timeout(global_events, ChunkGenerated, 1, condition=lambda x: isinstance(x.chunk, TokenChunk) and x.chunk.finish_reason is not None)
-
# TODO: The ideal with this test is if we had some tooling to scroll through the state, and say
# 'asser that there was a time that the error_type, error_message was not none and the failure count was nonzero'
# as we reset the failures back to zero when we have a successful inference.
- assert len(worker.assigned_runners[RUNNER_1_ID].failures) == 0
+ assert len(worker.assigned_runners[RUNNER_1_ID].failures) == 0
assert worker.state.tasks[TASK_1_ID].error_type is None
assert worker.state.tasks[TASK_1_ID].error_message is None
events = await global_events.get_events_since(0)
- assert len([x for x in events if isinstance(x.event, RunnerStatusUpdated) and isinstance(x.event.runner_status, FailedRunnerStatus)]) == 1
- assert len([x for x in events if isinstance(x.event, TaskStateUpdated) and x.event.task_status == TaskStatus.FAILED]) == 1
+ assert (
+ len(
+ [
+ x
+ for x in events
+ if isinstance(x.event, RunnerStatusUpdated)
+ and isinstance(x.event.runner_status, FailedRunnerStatus)
+ ]
+ )
+ == 1
+ )
+ assert (
+ len(
+ [
+ x
+ for x in events
+ if isinstance(x.event, TaskStateUpdated)
+ and x.event.task_status == TaskStatus.FAILED
+ ]
+ )
+ == 1
+ )
- response_string = ''
+ response_string = ""
events = await global_events.get_events_since(0)
seen_task_started, seen_task_finished = False, False
@@ -164,7 +218,7 @@ async def test_stream_response_failed_once(
assert isinstance(event.chunk, TokenChunk)
response_string += event.chunk.text
- assert 'queen' in response_string.lower()
+ assert "queen" in response_string.lower()
assert seen_task_started
assert seen_task_finished
@@ -173,17 +227,19 @@ async def test_stream_response_failed_once(
InstanceDeleted(
instance_id=instance_value.instance_id,
),
- ],
- origin=MASTER_NODE_ID
+ ],
+ origin=MASTER_NODE_ID,
)
await asyncio.sleep(0.3)
async def test_stream_response_timeout(
- worker_running: Callable[[NodeId], Awaitable[tuple[Worker, AsyncSQLiteEventStorage]]],
+ worker_running: Callable[
+ [NodeId], Awaitable[tuple[Worker, AsyncSQLiteEventStorage]]
+ ],
instance: Callable[[InstanceId, NodeId, RunnerId], Instance],
- chat_completion_task: Callable[[InstanceId, TaskId], Task]
+ chat_completion_task: Callable[[InstanceId, TaskId], Task],
):
_, global_events = await worker_running(NODE_A)
@@ -191,30 +247,60 @@ async def test_stream_response_timeout(
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'
+ task.task_params.messages[0].content = "EXO RUNNER MUST TIMEOUT"
await global_events.append_events(
[
InstanceCreated(instance=instance_value),
- TaskCreated(task_id=task.task_id, task=task)
- ],
- origin=MASTER_NODE_ID
+ TaskCreated(task_id=task.task_id, task=task),
+ ],
+ origin=MASTER_NODE_ID,
)
await until_event_with_timeout(global_events, TaskFailed, multiplicity=3)
events = await global_events.get_events_since(0)
print(events)
- assert len([x for x in events if isinstance(x.event, RunnerStatusUpdated) and isinstance(x.event.runner_status, FailedRunnerStatus)]) == 3
- assert len([x for x in events if isinstance(x.event, TaskStateUpdated) and x.event.task_status == TaskStatus.FAILED]) == 3
- assert len([x for x in events if isinstance(x.event, TaskFailed) and 'timeouterror' in x.event.error_type.lower()]) == 3
+ assert (
+ len(
+ [
+ x
+ for x in events
+ if isinstance(x.event, RunnerStatusUpdated)
+ and isinstance(x.event.runner_status, FailedRunnerStatus)
+ ]
+ )
+ == 3
+ )
+ assert (
+ len(
+ [
+ x
+ for x in events
+ if isinstance(x.event, TaskStateUpdated)
+ and x.event.task_status == TaskStatus.FAILED
+ ]
+ )
+ == 3
+ )
+ assert (
+ len(
+ [
+ x
+ for x in events
+ if isinstance(x.event, TaskFailed)
+ and "timeouterror" in x.event.error_type.lower()
+ ]
+ )
+ == 3
+ )
await global_events.append_events(
[
InstanceDeleted(
instance_id=instance_value.instance_id,
),
- ],
- origin=MASTER_NODE_ID
+ ],
+ origin=MASTER_NODE_ID,
)
- await asyncio.sleep(0.3)
\ No newline at end of file
+ await asyncio.sleep(0.3)
diff --git a/src/exo/worker/tests/test_integration/test_instantiation.py b/src/exo/worker/tests/test_integration/test_instantiation.py
index 21f296b1..dc0773b2 100644
--- a/src/exo/worker/tests/test_integration/test_instantiation.py
+++ b/src/exo/worker/tests/test_integration/test_instantiation.py
@@ -30,22 +30,21 @@ from exo.worker.tests.test_integration.integration_utils import until_event_with
async def test_runner_spinup_exception(
- worker_running: Callable[[NodeId], Awaitable[tuple[Worker, AsyncSQLiteEventStorage]]],
+ worker_running: Callable[
+ [NodeId], Awaitable[tuple[Worker, AsyncSQLiteEventStorage]]
+ ],
instance: Callable[[InstanceId, NodeId, RunnerId], Instance],
- ):
+):
_, global_events = await worker_running(NODE_A)
instance_value: Instance = instance(INSTANCE_1_ID, NODE_A, RUNNER_1_ID)
instance_value.instance_type = InstanceStatus.ACTIVE
- instance_value.shard_assignments.runner_to_shard[RUNNER_1_ID].immediate_exception = True
+ instance_value.shard_assignments.runner_to_shard[
+ RUNNER_1_ID
+ ].immediate_exception = True
await global_events.append_events(
- [
- InstanceCreated(
- instance=instance_value
- )
- ],
- origin=MASTER_NODE_ID
+ [InstanceCreated(instance=instance_value)], origin=MASTER_NODE_ID
)
await asyncio.sleep(5.0)
@@ -53,17 +52,28 @@ async def test_runner_spinup_exception(
# Ensure the correct events have been emitted
events = await global_events.get_events_since(0)
- assert len([x for x in events if isinstance(x.event, RunnerStatusUpdated) \
- and isinstance(x.event.runner_status, FailedRunnerStatus) \
- and x.event.runner_status.error_message is not None \
- and 'fake exception' in x.event.runner_status.error_message.lower()]) == 3
+ assert (
+ len(
+ [
+ x
+ for x in events
+ if isinstance(x.event, RunnerStatusUpdated)
+ and isinstance(x.event.runner_status, FailedRunnerStatus)
+ and x.event.runner_status.error_message is not None
+ and "fake exception" in x.event.runner_status.error_message.lower()
+ ]
+ )
+ == 3
+ )
assert any([isinstance(x.event, InstanceDeleted) for x in events])
async def test_runner_spinup_timeout(
- worker_running: Callable[[NodeId], Awaitable[tuple[Worker, AsyncSQLiteEventStorage]]],
+ worker_running: Callable[
+ [NodeId], Awaitable[tuple[Worker, AsyncSQLiteEventStorage]]
+ ],
instance: Callable[[InstanceId, NodeId, RunnerId], Instance],
- ):
+):
_, global_events = await worker_running(NODE_A)
instance_value: Instance = instance(INSTANCE_1_ID, NODE_A, RUNNER_1_ID)
@@ -71,18 +81,28 @@ async def test_runner_spinup_timeout(
instance_value.shard_assignments.runner_to_shard[RUNNER_1_ID].should_timeout = 10
await global_events.append_events(
- [
- InstanceCreated(
- instance=instance_value
- )
- ],
- origin=MASTER_NODE_ID
+ [InstanceCreated(instance=instance_value)], origin=MASTER_NODE_ID
)
- await until_event_with_timeout(global_events, RunnerStatusUpdated, multiplicity=3, condition=lambda x: isinstance(x.runner_status, FailedRunnerStatus))
+ await until_event_with_timeout(
+ global_events,
+ RunnerStatusUpdated,
+ multiplicity=3,
+ condition=lambda x: isinstance(x.runner_status, FailedRunnerStatus),
+ )
# Ensure the correct events have been emitted
events = await global_events.get_events_since(0)
- assert len([x for x in events if isinstance(x.event, RunnerStatusUpdated) and isinstance(x.event.runner_status, FailedRunnerStatus)]) == 3
- assert any([isinstance(x.event, InstanceDeleted) for x in events])
\ No newline at end of file
+ assert (
+ len(
+ [
+ x
+ for x in events
+ if isinstance(x.event, RunnerStatusUpdated)
+ and isinstance(x.event.runner_status, FailedRunnerStatus)
+ ]
+ )
+ == 3
+ )
+ assert any([isinstance(x.event, InstanceDeleted) for x in events])
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 a84b52d5..beb73acf 100644
--- a/src/exo/worker/tests/test_integration/test_instantiation_sad.py
+++ b/src/exo/worker/tests/test_integration/test_instantiation_sad.py
@@ -30,22 +30,21 @@ from exo.worker.tests.test_integration.integration_utils import until_event_with
async def test_runner_spinup_exception(
- worker_running: Callable[[NodeId], Awaitable[tuple[Worker, AsyncSQLiteEventStorage]]],
+ worker_running: Callable[
+ [NodeId], Awaitable[tuple[Worker, AsyncSQLiteEventStorage]]
+ ],
instance: Callable[[InstanceId, NodeId, RunnerId], Instance],
- ):
+):
_, global_events = await worker_running(NODE_A)
instance_value: Instance = instance(INSTANCE_1_ID, NODE_A, RUNNER_1_ID)
instance_value.instance_type = InstanceStatus.ACTIVE
- instance_value.shard_assignments.runner_to_shard[RUNNER_1_ID].immediate_exception = True
+ instance_value.shard_assignments.runner_to_shard[
+ RUNNER_1_ID
+ ].immediate_exception = True
await global_events.append_events(
- [
- InstanceCreated(
- instance=instance_value
- )
- ],
- origin=MASTER_NODE_ID
+ [InstanceCreated(instance=instance_value)], origin=MASTER_NODE_ID
)
await asyncio.sleep(5.0)
@@ -53,14 +52,26 @@ async def test_runner_spinup_exception(
# Ensure the correct events have been emitted
events = await global_events.get_events_since(0)
- assert len([x for x in events if isinstance(x.event, RunnerStatusUpdated) and isinstance(x.event.runner_status, FailedRunnerStatus)]) == 3
+ assert (
+ len(
+ [
+ x
+ for x in events
+ if isinstance(x.event, RunnerStatusUpdated)
+ and isinstance(x.event.runner_status, FailedRunnerStatus)
+ ]
+ )
+ == 3
+ )
assert any([isinstance(x.event, InstanceDeleted) for x in events])
async def test_runner_spinup_timeout(
- worker_running: Callable[[NodeId], Awaitable[tuple[Worker, AsyncSQLiteEventStorage]]],
+ worker_running: Callable[
+ [NodeId], Awaitable[tuple[Worker, AsyncSQLiteEventStorage]]
+ ],
instance: Callable[[InstanceId, NodeId, RunnerId], Instance],
- ):
+):
_, global_events = await worker_running(NODE_A)
instance_value: Instance = instance(INSTANCE_1_ID, NODE_A, RUNNER_1_ID)
@@ -68,18 +79,28 @@ async def test_runner_spinup_timeout(
instance_value.shard_assignments.runner_to_shard[RUNNER_1_ID].should_timeout = 10
await global_events.append_events(
- [
- InstanceCreated(
- instance=instance_value
- )
- ],
- origin=MASTER_NODE_ID
+ [InstanceCreated(instance=instance_value)], origin=MASTER_NODE_ID
)
- await until_event_with_timeout(global_events, RunnerStatusUpdated, multiplicity=3, condition=lambda x: isinstance(x.runner_status, FailedRunnerStatus))
+ await until_event_with_timeout(
+ global_events,
+ RunnerStatusUpdated,
+ multiplicity=3,
+ condition=lambda x: isinstance(x.runner_status, FailedRunnerStatus),
+ )
# Ensure the correct events have been emitted
events = await global_events.get_events_since(0)
- assert len([x for x in events if isinstance(x.event, RunnerStatusUpdated) and isinstance(x.event.runner_status, FailedRunnerStatus)]) == 3
- assert any([isinstance(x.event, InstanceDeleted) for x in events])
\ No newline at end of file
+ assert (
+ len(
+ [
+ x
+ for x in events
+ if isinstance(x.event, RunnerStatusUpdated)
+ and isinstance(x.event.runner_status, FailedRunnerStatus)
+ ]
+ )
+ == 3
+ )
+ assert any([isinstance(x.event, InstanceDeleted) for x in events])
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 c6c96197..b38ab16d 100644
--- a/src/exo/worker/tests/test_multimodel/test_inference_llama70B.py
+++ b/src/exo/worker/tests/test_multimodel/test_inference_llama70B.py
@@ -16,7 +16,13 @@ from exo.shared.types.events import (
TaskCreated,
)
from exo.shared.types.models import ModelId, ModelMetadata
-from exo.shared.types.tasks import ChatCompletionTask, Task, TaskId, TaskStatus, TaskType
+from exo.shared.types.tasks import (
+ ChatCompletionTask,
+ Task,
+ TaskId,
+ TaskStatus,
+ TaskType,
+)
from exo.shared.types.worker.common import InstanceId
from exo.shared.types.worker.instances import (
Instance,
@@ -43,14 +49,14 @@ from exo.worker.tests.test_integration.integration_utils import (
)
from exo.worker.worker import Worker
-MODEL_ID = 'mlx-community/Llama-3.3-70B-Instruct-4bit'
+MODEL_ID = "mlx-community/Llama-3.3-70B-Instruct-4bit"
+
@pytest.fixture
async def model_meta() -> ModelMetadata:
return await get_model_meta(MODEL_ID)
-
def _get_model_size_gb(path: str) -> float:
"""Calculate total size of directory recursively in GB."""
total_size = 0
@@ -61,19 +67,29 @@ def _get_model_size_gb(path: str) -> float:
total_size += os.path.getsize(filepath)
return total_size / (1024**3) # Convert bytes to GB
+
@pytest.mark.skipif(
not (
- os.path.exists(os.path.expanduser("~/.exo/models/mlx-community--Llama-3.3-70B-Instruct-4bit/"))
- and _get_model_size_gb(os.path.expanduser("~/.exo/models/mlx-community--Llama-3.3-70B-Instruct-4bit/")) > 30
+ os.path.exists(
+ os.path.expanduser(
+ "~/.exo/models/mlx-community--Llama-3.3-70B-Instruct-4bit/"
+ )
+ )
+ and _get_model_size_gb(
+ os.path.expanduser(
+ "~/.exo/models/mlx-community--Llama-3.3-70B-Instruct-4bit/"
+ )
+ )
+ > 30
),
- reason="This test only runs when model mlx-community/Llama-3.3-70B-Instruct-4bit is downloaded"
+ reason="This test only runs when model mlx-community/Llama-3.3-70B-Instruct-4bit is downloaded",
)
async def test_2_runner_inference(
logger: Logger,
pipeline_shard_meta: Callable[[int, int], PipelineShardMetadata],
hosts: Callable[[int], list[Host]],
- chat_completion_task: Callable[[InstanceId, TaskId], Task]
- ):
+ chat_completion_task: Callable[[InstanceId, TaskId], Task],
+):
event_log_manager = EventLogManager(EventLogConfig(), logger)
await event_log_manager.initialize()
shard_downloader = NoopShardDownloader()
@@ -81,10 +97,22 @@ async def test_2_runner_inference(
global_events = event_log_manager.global_events
await global_events.delete_all_events()
- worker1 = Worker(NODE_A, logger=logger, shard_downloader=shard_downloader, worker_events=global_events, global_events=global_events)
+ worker1 = Worker(
+ NODE_A,
+ logger=logger,
+ shard_downloader=shard_downloader,
+ worker_events=global_events,
+ global_events=global_events,
+ )
asyncio.create_task(run(worker1, logger))
- worker2 = Worker(NODE_B, logger=logger, shard_downloader=shard_downloader, worker_events=global_events, global_events=global_events)
+ worker2 = Worker(
+ NODE_B,
+ logger=logger,
+ shard_downloader=shard_downloader,
+ worker_events=global_events,
+ global_events=global_events,
+ )
asyncio.create_task(run(worker2, logger))
## Instance
@@ -94,44 +122,43 @@ async def test_2_runner_inference(
model_id=model_id,
runner_to_shard={
RUNNER_1_ID: pipeline_shard_meta(2, 0),
- RUNNER_2_ID: pipeline_shard_meta(2, 1)
+ RUNNER_2_ID: pipeline_shard_meta(2, 1),
},
- node_to_runner={
- NODE_A: RUNNER_1_ID,
- NODE_B: RUNNER_2_ID
- }
+ node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID},
)
-
+
instance = Instance(
instance_id=INSTANCE_1_ID,
instance_type=InstanceStatus.ACTIVE,
shard_assignments=shard_assignments,
- hosts=hosts(2)
+ hosts=hosts(2),
)
task = chat_completion_task(INSTANCE_1_ID, TASK_1_ID)
- task.task_params.messages[0].content = 'Can you explain to me how a bubble sort works, speaking as if you are a fairy.'
+ task.task_params.messages[
+ 0
+ ].content = (
+ "Can you explain to me how a bubble sort works, speaking as if you are a fairy."
+ )
task.task_params.max_tokens = 1000
await global_events.append_events(
[
- InstanceCreated(
- instance=instance
- ),
- TaskCreated(
- task_id=task.task_id,
- task=task
- )
- ],
- origin=MASTER_NODE_ID
+ InstanceCreated(instance=instance),
+ TaskCreated(task_id=task.task_id, task=task),
+ ],
+ origin=MASTER_NODE_ID,
)
- seen_task_started, seen_task_finished, response_string = await read_streaming_response(global_events)
+ (
+ seen_task_started,
+ seen_task_finished,
+ response_string,
+ ) = await read_streaming_response(global_events)
assert seen_task_started
assert seen_task_finished
- assert 'swap' in response_string.lower()
-
+ assert "swap" in response_string.lower()
idx = await global_events.get_last_idx()
await asyncio.sleep(1.0)
@@ -143,27 +170,35 @@ async def test_2_runner_inference(
InstanceDeleted(
instance_id=instance.instance_id,
),
- ],
- origin=MASTER_NODE_ID
+ ],
+ origin=MASTER_NODE_ID,
)
await asyncio.sleep(2.0)
-
@pytest.mark.skipif(
not (
- os.path.exists(os.path.expanduser("~/.exo/models/mlx-community--Llama-3.3-70B-Instruct-4bit/"))
- and _get_model_size_gb(os.path.expanduser("~/.exo/models/mlx-community--Llama-3.3-70B-Instruct-4bit/")) > 30
+ os.path.exists(
+ os.path.expanduser(
+ "~/.exo/models/mlx-community--Llama-3.3-70B-Instruct-4bit/"
+ )
+ )
+ and _get_model_size_gb(
+ os.path.expanduser(
+ "~/.exo/models/mlx-community--Llama-3.3-70B-Instruct-4bit/"
+ )
+ )
+ > 30
),
- reason="This test only runs when model mlx-community/Llama-3.3-70B-Instruct-4bit is downloaded"
+ reason="This test only runs when model mlx-community/Llama-3.3-70B-Instruct-4bit is downloaded",
)
async def test_parallel_inference(
logger: Logger,
pipeline_shard_meta: Callable[[int, int], PipelineShardMetadata],
hosts: Callable[[int], list[Host]],
- chat_completion_task: Callable[[InstanceId, TaskId], Task]
- ):
+ chat_completion_task: Callable[[InstanceId, TaskId], Task],
+):
event_log_manager = EventLogManager(EventLogConfig(), logger)
await event_log_manager.initialize()
shard_downloader = NoopShardDownloader()
@@ -171,10 +206,22 @@ async def test_parallel_inference(
global_events = event_log_manager.global_events
await global_events.delete_all_events()
- worker1 = Worker(NODE_A, logger=logger, shard_downloader=shard_downloader, worker_events=global_events, global_events=global_events)
+ worker1 = Worker(
+ NODE_A,
+ logger=logger,
+ shard_downloader=shard_downloader,
+ worker_events=global_events,
+ global_events=global_events,
+ )
asyncio.create_task(run(worker1, logger))
- worker2 = Worker(NODE_B, logger=logger, shard_downloader=shard_downloader, worker_events=global_events, global_events=global_events)
+ worker2 = Worker(
+ NODE_B,
+ logger=logger,
+ shard_downloader=shard_downloader,
+ worker_events=global_events,
+ global_events=global_events,
+ )
asyncio.create_task(run(worker2, logger))
## Instance
@@ -184,26 +231,27 @@ async def test_parallel_inference(
model_id=model_id,
runner_to_shard={
RUNNER_1_ID: pipeline_shard_meta(2, 0),
- RUNNER_2_ID: pipeline_shard_meta(2, 1)
+ RUNNER_2_ID: pipeline_shard_meta(2, 1),
},
- node_to_runner={
- NODE_A: RUNNER_1_ID,
- NODE_B: RUNNER_2_ID
- }
+ node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID},
)
-
+
instance = Instance(
instance_id=INSTANCE_1_ID,
instance_type=InstanceStatus.ACTIVE,
shard_assignments=shard_assignments,
- hosts=hosts(2)
+ hosts=hosts(2),
)
completion_create_params_1 = ChatCompletionTaskParams(
model="gpt-4",
- messages=[ChatCompletionMessage(role="user", content='Tell me a haiku that uses the word "pond".')],
+ messages=[
+ ChatCompletionMessage(
+ role="user", content='Tell me a haiku that uses the word "pond".'
+ )
+ ],
stream=True,
- max_tokens=1000
+ max_tokens=1000,
)
task1 = ChatCompletionTask(
task_id=TASK_1_ID,
@@ -211,14 +259,18 @@ async def test_parallel_inference(
instance_id=INSTANCE_1_ID,
task_type=TaskType.CHAT_COMPLETION,
task_status=TaskStatus.PENDING,
- task_params=completion_create_params_1
+ task_params=completion_create_params_1,
)
-
+
completion_create_params_2 = ChatCompletionTaskParams(
model="gpt-4",
- messages=[ChatCompletionMessage(role="user", content='Tell me a haiku that uses the word "tree".')],
+ messages=[
+ ChatCompletionMessage(
+ role="user", content='Tell me a haiku that uses the word "tree".'
+ )
+ ],
stream=True,
- max_tokens=1000
+ max_tokens=1000,
)
task2 = ChatCompletionTask(
task_id=TASK_2_ID,
@@ -226,30 +278,34 @@ async def test_parallel_inference(
instance_id=INSTANCE_1_ID,
task_type=TaskType.CHAT_COMPLETION,
task_status=TaskStatus.PENDING,
- task_params=completion_create_params_2
- )
+ task_params=completion_create_params_2,
+ )
await global_events.append_events(
[
- InstanceCreated(
- instance=instance
- ),
- TaskCreated(
- task_id=task1.task_id,
- task=task1
- ),
- TaskCreated(
- task_id=task2.task_id,
- task=task2
- ),
- ],
- origin=MASTER_NODE_ID
+ InstanceCreated(instance=instance),
+ TaskCreated(task_id=task1.task_id, task=task1),
+ TaskCreated(task_id=task2.task_id, task=task2),
+ ],
+ origin=MASTER_NODE_ID,
)
- seen_task_started_1, seen_task_finished_1, response_string_1 = await read_streaming_response(global_events)
+ (
+ seen_task_started_1,
+ seen_task_finished_1,
+ response_string_1,
+ ) = await read_streaming_response(global_events)
- incomplete_task = TASK_2_ID if worker1.state.tasks[TASK_1_ID].task_status == TaskStatus.COMPLETE else TASK_2_ID
- seen_task_started_2, seen_task_finished_2, response_string_2 = await read_streaming_response(global_events, filter_task=incomplete_task)
+ incomplete_task = (
+ TASK_2_ID
+ if worker1.state.tasks[TASK_1_ID].task_status == TaskStatus.COMPLETE
+ else TASK_2_ID
+ )
+ (
+ seen_task_started_2,
+ seen_task_finished_2,
+ response_string_2,
+ ) = await read_streaming_response(global_events, filter_task=incomplete_task)
assert seen_task_started_1
assert seen_task_finished_1
@@ -259,14 +315,13 @@ async def test_parallel_inference(
print(response_string_1)
print(response_string_2)
- assert (
- ('pond' in response_string_1.lower()) ^ ('pond' in response_string_2.lower())
+ assert ("pond" in response_string_1.lower()) ^ (
+ "pond" in response_string_2.lower()
), "'pond' must appear in exactly one response"
- assert (
- ('tree' in response_string_1.lower()) ^ ('tree' in response_string_2.lower())
+ assert ("tree" in response_string_1.lower()) ^ (
+ "tree" in response_string_2.lower()
), "'tree' must appear in exactly one response"
-
idx = await global_events.get_last_idx()
await asyncio.sleep(1.0)
events = await global_events.get_events_since(idx)
@@ -277,8 +332,8 @@ async def test_parallel_inference(
InstanceDeleted(
instance_id=instance.instance_id,
),
- ],
- origin=MASTER_NODE_ID
+ ],
+ origin=MASTER_NODE_ID,
)
- await asyncio.sleep(2.0)
\ No newline at end of file
+ await asyncio.sleep(2.0)
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 d6ae4e7c..dd304bd1 100644
--- a/src/exo/worker/tests/test_plan/test_worker_plan.py
+++ b/src/exo/worker/tests/test_plan/test_worker_plan.py
@@ -72,38 +72,41 @@ def _get_test_cases() -> list[PlanTestCase]:
PlanTestCase(
description="no runners -> no-op",
in_process_runners=[],
- state=State(node_status={NODE_A: NodeStatus.Idle}, instances={}, runners={}),
+ state=State(
+ node_status={NODE_A: NodeStatus.Idle}, instances={}, runners={}
+ ),
expected_op=None,
),
-
# Both 'assigned' and 'downloading' should be blocking ops - so if we are in either of these we should unassign to retry.
# This needs to change when we move to an async worker
make_test_case(
description="runner state assigned, runner is assigned and downloading -> unassign",
- runner_specs=[{
- 'runner_id': RUNNER_1_ID,
- 'node_id': NODE_A,
- 'device_rank': 0,
- 'status': make_downloading_status(NODE_A),
- 'downloaded': False
- }],
+ runner_specs=[
+ {
+ "runner_id": RUNNER_1_ID,
+ "node_id": NODE_A,
+ "device_rank": 0,
+ "status": make_downloading_status(NODE_A),
+ "downloaded": False,
+ }
+ ],
instance_status=InstanceStatus.INACTIVE,
expected_op=UnassignRunnerOp(runner_id=RUNNER_1_ID),
),
-
make_test_case(
description="ready runner, model present -> no-op",
- runner_specs=[{
- 'runner_id': RUNNER_1_ID,
- 'node_id': NODE_A,
- 'device_rank': 0,
- 'status': InactiveRunnerStatus(),
- 'downloaded': True
- }],
+ runner_specs=[
+ {
+ "runner_id": RUNNER_1_ID,
+ "node_id": NODE_A,
+ "device_rank": 0,
+ "status": InactiveRunnerStatus(),
+ "downloaded": True,
+ }
+ ],
instance_status=InstanceStatus.INACTIVE,
expected_op=None,
),
-
PlanTestCase(
description="runner assigned and not in state -> AssignRunnerOp",
in_process_runners=[],
@@ -125,10 +128,9 @@ def _get_test_cases() -> list[PlanTestCase]:
end_layer=1,
n_layers=1,
),
- hosts=[]
+ hosts=[],
),
),
-
PlanTestCase(
description="runner assigned but no longer in state -> UnassignRunnerOp",
in_process_runners=[
@@ -140,187 +142,206 @@ def _get_test_cases() -> list[PlanTestCase]:
downloaded=False,
)
],
- state=State(node_status={NODE_A: NodeStatus.Idle}, instances={}, runners={}),
+ state=State(
+ node_status={NODE_A: NodeStatus.Idle}, instances={}, runners={}
+ ),
expected_op=UnassignRunnerOp(runner_id=RUNNER_1_ID),
),
-
make_test_case(
description="ready runner (and state up) -> expect RunnerUpOp",
- runner_specs=[{
- 'runner_id': RUNNER_1_ID,
- 'node_id': NODE_A,
- 'device_rank': 0,
- 'status': InactiveRunnerStatus(),
- 'downloaded': True
- }],
+ runner_specs=[
+ {
+ "runner_id": RUNNER_1_ID,
+ "node_id": NODE_A,
+ "device_rank": 0,
+ "status": InactiveRunnerStatus(),
+ "downloaded": True,
+ }
+ ],
instance_status=InstanceStatus.ACTIVE,
expected_op=RunnerUpOp(runner_id=RUNNER_1_ID),
),
-
make_test_case(
description="1 ready, 1 downloading (and state up) -> no-op",
runner_specs=[
{
- 'runner_id': RUNNER_1_ID,
- 'node_id': NODE_A,
- 'device_rank': 0,
- 'status': InactiveRunnerStatus(),
- 'downloaded': True
+ "runner_id": RUNNER_1_ID,
+ "node_id": NODE_A,
+ "device_rank": 0,
+ "status": InactiveRunnerStatus(),
+ "downloaded": True,
},
{
- 'runner_id': RUNNER_2_ID,
- 'node_id': NODE_B,
- 'device_rank': 1,
- 'status': DownloadingRunnerStatus(download_progress=DownloadPending(node_id=NODE_A)),
- 'downloaded': False
+ "runner_id": RUNNER_2_ID,
+ "node_id": NODE_B,
+ "device_rank": 1,
+ "status": DownloadingRunnerStatus(
+ download_progress=DownloadPending(node_id=NODE_A)
+ ),
+ "downloaded": False,
+ },
+ ],
+ tasks=[
+ {
+ "task_id": TASK_1_ID,
+ "instance_id": INSTANCE_1_ID,
+ "status": TaskStatus.PENDING,
+ "messages": [{"role": "user", "content": "Hello, world!"}],
}
],
- tasks=[{
- 'task_id': TASK_1_ID,
- 'instance_id': INSTANCE_1_ID,
- 'status': TaskStatus.PENDING,
- 'messages': [{'role': 'user', 'content': 'Hello, world!'}]
- }],
instance_status=InstanceStatus.ACTIVE,
- expected_op=None
+ expected_op=None,
),
-
make_test_case(
description="2 ready runners (and state up) -> expect RunnerUpOp",
runner_specs=[
{
- 'runner_id': RUNNER_1_ID,
- 'node_id': NODE_A,
- 'device_rank': 0,
- 'status': InactiveRunnerStatus(),
- 'downloaded': True
+ "runner_id": RUNNER_1_ID,
+ "node_id": NODE_A,
+ "device_rank": 0,
+ "status": InactiveRunnerStatus(),
+ "downloaded": True,
},
{
- 'runner_id': RUNNER_2_ID,
- 'node_id': NODE_B,
- 'device_rank': 1,
- 'status': InactiveRunnerStatus(),
- 'downloaded': True
+ "runner_id": RUNNER_2_ID,
+ "node_id": NODE_B,
+ "device_rank": 1,
+ "status": InactiveRunnerStatus(),
+ "downloaded": True,
+ },
+ ],
+ tasks=[
+ {
+ "task_id": TASK_1_ID,
+ "instance_id": INSTANCE_1_ID,
+ "status": TaskStatus.PENDING,
+ "messages": [{"role": "user", "content": "Hello, world!"}],
}
],
- tasks=[{
- 'task_id': TASK_1_ID,
- 'instance_id': INSTANCE_1_ID,
- 'status': TaskStatus.PENDING,
- 'messages': [{'role': 'user', 'content': 'Hello, world!'}]
- }],
instance_status=InstanceStatus.ACTIVE,
- expected_op=RunnerUpOp(runner_id=RUNNER_1_ID)
+ expected_op=RunnerUpOp(runner_id=RUNNER_1_ID),
),
-
make_test_case(
description="loaded runner (and state down) -> expect RunnerDownOp",
- runner_specs=[{
- 'runner_id': RUNNER_1_ID,
- 'node_id': NODE_A,
- 'device_rank': 0,
- 'status': LoadedRunnerStatus(),
- 'downloaded': True
- }],
+ runner_specs=[
+ {
+ "runner_id": RUNNER_1_ID,
+ "node_id": NODE_A,
+ "device_rank": 0,
+ "status": LoadedRunnerStatus(),
+ "downloaded": True,
+ }
+ ],
instance_status=InstanceStatus.INACTIVE,
expected_op=RunnerDownOp(runner_id=RUNNER_1_ID),
),
-
make_test_case(
description="failed runner (and state down) -> expect RunnerDownOp",
- runner_specs=[{
- 'runner_id': RUNNER_1_ID,
- 'node_id': NODE_A,
- 'device_rank': 0,
- 'status': FailedRunnerStatus(),
- 'downloaded': True
- }],
+ runner_specs=[
+ {
+ "runner_id": RUNNER_1_ID,
+ "node_id": NODE_A,
+ "device_rank": 0,
+ "status": FailedRunnerStatus(),
+ "downloaded": True,
+ }
+ ],
instance_status=InstanceStatus.INACTIVE,
expected_op=RunnerDownOp(runner_id=RUNNER_1_ID),
),
-
make_test_case(
description="loaded runner, model present, task pending -> expect ExecuteTaskOp",
- runner_specs=[{
- 'runner_id': RUNNER_1_ID,
- 'node_id': NODE_A,
- 'device_rank': 0,
- 'status': LoadedRunnerStatus(),
- 'downloaded': True
- }],
- tasks=[{
- 'task_id': TASK_1_ID,
- 'instance_id': INSTANCE_1_ID,
- 'status': TaskStatus.PENDING,
- 'messages': [{'role': 'user', 'content': 'Hello, world!'}]
- }],
+ runner_specs=[
+ {
+ "runner_id": RUNNER_1_ID,
+ "node_id": NODE_A,
+ "device_rank": 0,
+ "status": LoadedRunnerStatus(),
+ "downloaded": True,
+ }
+ ],
+ tasks=[
+ {
+ "task_id": TASK_1_ID,
+ "instance_id": INSTANCE_1_ID,
+ "status": TaskStatus.PENDING,
+ "messages": [{"role": "user", "content": "Hello, world!"}],
+ }
+ ],
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_params=ChatCompletionTaskParams(
- model=str(MODEL_A_ID),
- messages=[ChatCompletionMessage(role="user", content="Hello, world!")]
+ 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_params=ChatCompletionTaskParams(
+ model=str(MODEL_A_ID),
+ messages=[
+ ChatCompletionMessage(role="user", content="Hello, world!")
+ ],
+ ),
),
- )),
+ ),
),
-
# We should only run rank 0 once all other ranks are running.
make_test_case(
description="two loaded runners & task, i'm rank 0 -> no-op",
runner_specs=[
{
- 'runner_id': RUNNER_1_ID,
- 'node_id': NODE_A,
- 'device_rank': 0,
- 'status': LoadedRunnerStatus(),
- 'downloaded': True
+ "runner_id": RUNNER_1_ID,
+ "node_id": NODE_A,
+ "device_rank": 0,
+ "status": LoadedRunnerStatus(),
+ "downloaded": True,
},
{
- 'runner_id': RUNNER_2_ID,
- 'node_id': NODE_B,
- 'device_rank': 1,
- 'status': LoadedRunnerStatus(),
- 'downloaded': True
+ "runner_id": RUNNER_2_ID,
+ "node_id": NODE_B,
+ "device_rank": 1,
+ "status": LoadedRunnerStatus(),
+ "downloaded": True,
+ },
+ ],
+ tasks=[
+ {
+ "task_id": TASK_1_ID,
+ "instance_id": INSTANCE_1_ID,
+ "status": TaskStatus.PENDING,
+ "messages": [{"role": "user", "content": "Hello, world!"}],
}
],
- tasks=[{
- 'task_id': TASK_1_ID,
- 'instance_id': INSTANCE_1_ID,
- 'status': TaskStatus.PENDING,
- 'messages': [{'role': 'user', 'content': 'Hello, world!'}]
- }],
instance_status=InstanceStatus.ACTIVE,
- expected_op=None
+ expected_op=None,
),
-
make_test_case(
description="two loaded runners & task, i'm rank 1 -> expect ExecuteTaskOp on rank 1",
runner_specs=[
{
- 'runner_id': RUNNER_1_ID,
- 'node_id': NODE_A,
- 'device_rank': 1,
- 'status': LoadedRunnerStatus(),
- 'downloaded': True
+ "runner_id": RUNNER_1_ID,
+ "node_id": NODE_A,
+ "device_rank": 1,
+ "status": LoadedRunnerStatus(),
+ "downloaded": True,
+ },
+ {
+ "runner_id": RUNNER_2_ID,
+ "node_id": NODE_B,
+ "device_rank": 0,
+ "status": LoadedRunnerStatus(),
+ "downloaded": True,
},
+ ],
+ tasks=[
{
- 'runner_id': RUNNER_2_ID,
- 'node_id': NODE_B,
- 'device_rank': 0,
- 'status': LoadedRunnerStatus(),
- 'downloaded': True
+ "task_id": TASK_1_ID,
+ "instance_id": INSTANCE_1_ID,
+ "status": TaskStatus.PENDING,
+ "messages": [{"role": "user", "content": "Hello, world!"}],
}
],
- tasks=[{
- 'task_id': TASK_1_ID,
- 'instance_id': INSTANCE_1_ID,
- 'status': TaskStatus.PENDING,
- 'messages': [{'role': 'user', 'content': 'Hello, world!'}]
- }],
instance_status=InstanceStatus.ACTIVE,
expected_op=ExecuteTaskOp(
runner_id=RUNNER_1_ID,
@@ -331,37 +352,40 @@ def _get_test_cases() -> list[PlanTestCase]:
task_type=TaskType.CHAT_COMPLETION,
task_params=ChatCompletionTaskParams(
model=str(MODEL_A_ID),
- messages=[ChatCompletionMessage(role="user", content="Hello, world!")],
+ messages=[
+ ChatCompletionMessage(role="user", content="Hello, world!")
+ ],
),
task_status=TaskStatus.PENDING,
),
),
),
-
make_test_case(
description="rank 1 loaded, rank 0 ready, i'm rank 0 -> expect ExecuteTaskOp on rank 0",
runner_specs=[
{
- 'runner_id': RUNNER_1_ID,
- 'node_id': NODE_A,
- 'device_rank': 0,
- 'status': LoadedRunnerStatus(),
- 'downloaded': True
+ "runner_id": RUNNER_1_ID,
+ "node_id": NODE_A,
+ "device_rank": 0,
+ "status": LoadedRunnerStatus(),
+ "downloaded": True,
+ },
+ {
+ "runner_id": RUNNER_2_ID,
+ "node_id": NODE_B,
+ "device_rank": 1,
+ "status": RunningRunnerStatus(),
+ "downloaded": True,
},
+ ],
+ tasks=[
{
- 'runner_id': RUNNER_2_ID,
- 'node_id': NODE_B,
- 'device_rank': 1,
- 'status': RunningRunnerStatus(),
- 'downloaded': True
+ "task_id": TASK_1_ID,
+ "instance_id": INSTANCE_1_ID,
+ "status": TaskStatus.PENDING,
+ "messages": [{"role": "user", "content": "Hello, world!"}],
}
],
- tasks=[{
- 'task_id': TASK_1_ID,
- 'instance_id': INSTANCE_1_ID,
- 'status': TaskStatus.PENDING,
- 'messages': [{'role': 'user', 'content': 'Hello, world!'}]
- }],
instance_status=InstanceStatus.ACTIVE,
expected_op=ExecuteTaskOp(
runner_id=RUNNER_1_ID,
@@ -372,93 +396,91 @@ def _get_test_cases() -> list[PlanTestCase]:
task_type=TaskType.CHAT_COMPLETION,
task_params=ChatCompletionTaskParams(
model=str(MODEL_A_ID),
- messages=[ChatCompletionMessage(role="user", content="Hello, world!")],
+ messages=[
+ ChatCompletionMessage(role="user", content="Hello, world!")
+ ],
),
task_status=TaskStatus.PENDING,
),
),
),
-
make_test_case(
description="this runner failed (1 node) -> RunnerDownOp",
- runner_specs=[{
- 'runner_id': RUNNER_1_ID,
- 'node_id': NODE_A,
- 'device_rank': 0,
- 'status': FailedRunnerStatus(),
- 'downloaded': True
- }],
+ runner_specs=[
+ {
+ "runner_id": RUNNER_1_ID,
+ "node_id": NODE_A,
+ "device_rank": 0,
+ "status": FailedRunnerStatus(),
+ "downloaded": True,
+ }
+ ],
instance_status=InstanceStatus.ACTIVE,
- expected_op=RunnerDownOp(runner_id=RUNNER_1_ID)
+ expected_op=RunnerDownOp(runner_id=RUNNER_1_ID),
),
-
make_test_case(
description="other runner failed -> RunnerDownOp",
runner_specs=[
{
- 'runner_id': RUNNER_1_ID,
- 'node_id': NODE_A,
- 'device_rank': 0,
- 'status': LoadedRunnerStatus(),
- 'downloaded': True
+ "runner_id": RUNNER_1_ID,
+ "node_id": NODE_A,
+ "device_rank": 0,
+ "status": LoadedRunnerStatus(),
+ "downloaded": True,
},
{
- 'runner_id': RUNNER_2_ID,
- 'node_id': NODE_B,
- 'device_rank': 1,
- 'status': FailedRunnerStatus(),
- 'downloaded': True
- }
+ "runner_id": RUNNER_2_ID,
+ "node_id": NODE_B,
+ "device_rank": 1,
+ "status": FailedRunnerStatus(),
+ "downloaded": True,
+ },
],
instance_status=InstanceStatus.ACTIVE,
- expected_op=RunnerDownOp(runner_id=RUNNER_1_ID)
+ expected_op=RunnerDownOp(runner_id=RUNNER_1_ID),
),
-
-
make_test_case(
description="this runner failed (2 nodes) -> no-op",
runner_specs=[
{
- 'runner_id': RUNNER_1_ID,
- 'node_id': NODE_A,
- 'device_rank': 0,
- 'status': FailedRunnerStatus(),
- 'downloaded': True
+ "runner_id": RUNNER_1_ID,
+ "node_id": NODE_A,
+ "device_rank": 0,
+ "status": FailedRunnerStatus(),
+ "downloaded": True,
},
{
- 'runner_id': RUNNER_2_ID,
- 'node_id': NODE_B,
- 'device_rank': 1,
- 'status': LoadedRunnerStatus(),
- 'downloaded': True
- }
+ "runner_id": RUNNER_2_ID,
+ "node_id": NODE_B,
+ "device_rank": 1,
+ "status": LoadedRunnerStatus(),
+ "downloaded": True,
+ },
],
instance_status=InstanceStatus.ACTIVE,
- expected_op=None
+ expected_op=None,
),
-
make_test_case(
description="this node failed, other node spun down -> RunnerDownOp",
runner_specs=[
{
- 'runner_id': RUNNER_1_ID,
- 'node_id': NODE_A,
- 'device_rank': 0,
- 'status': FailedRunnerStatus(),
- 'downloaded': True
+ "runner_id": RUNNER_1_ID,
+ "node_id": NODE_A,
+ "device_rank": 0,
+ "status": FailedRunnerStatus(),
+ "downloaded": True,
},
{
- 'runner_id': RUNNER_2_ID,
- 'node_id': NODE_B,
- 'device_rank': 1,
- 'status': InactiveRunnerStatus(),
- 'downloaded': True
- }
+ "runner_id": RUNNER_2_ID,
+ "node_id": NODE_B,
+ "device_rank": 1,
+ "status": InactiveRunnerStatus(),
+ "downloaded": True,
+ },
],
instance_status=InstanceStatus.ACTIVE,
- expected_op=RunnerDownOp(runner_id=RUNNER_1_ID)
+ expected_op=RunnerDownOp(runner_id=RUNNER_1_ID),
),
-
]
@@ -486,40 +508,46 @@ def test_worker_plan(case: PlanTestCase) -> None:
logger = logging.getLogger("test_worker_plan")
shard_downloader = NoopShardDownloader()
- worker = Worker(node_id=node_id, shard_downloader=shard_downloader, worker_events=None, global_events=None, logger=logger)
+ worker = Worker(
+ node_id=node_id,
+ shard_downloader=shard_downloader,
+ worker_events=None,
+ global_events=None,
+ logger=logger,
+ )
runner_config: InProcessRunner
for runner_config in case.in_process_runners:
-
- if len(case.state.instances) == 1:
+ if len(case.state.instances) == 1:
instance_id = next(iter(case.state.instances))
shard_assignments = case.state.instances[instance_id].shard_assignments
shard_metadata = shard_assignments.runner_to_shard[runner_config.runner_id]
-
+
# Only add this runner if it belongs to our node
runner_node = None
for node, runner in shard_assignments.node_to_runner.items():
if runner == runner_config.runner_id:
runner_node = node
break
-
+
if runner_node != node_id:
# This runner belongs to a different node, skip it
continue
-
+
elif len(case.state.instances) == 0:
shard_metadata = PipelineShardMetadata(
device_rank=runner_config.device_rank,
- world_size=1,
+ world_size=1,
model_meta=make_model_meta(runner_config.model_id),
start_layer=0,
end_layer=1,
n_layers=1,
)
else:
- raise Exception('test_worker_plan not currently designed to have more than 1 instance.')
-
+ raise Exception(
+ "test_worker_plan not currently designed to have more than 1 instance."
+ )
assigned_runner = AssignedRunner(
runner_id=runner_config.runner_id,
@@ -531,10 +559,11 @@ def test_worker_plan(case: PlanTestCase) -> None:
)
worker.assigned_runners[runner_config.runner_id] = assigned_runner
- op = plan(worker.assigned_runners,
- NODE_A,
- case.state.instances,
- case.state.runners,
- case.state.tasks,
- )
+ op = plan(
+ worker.assigned_runners,
+ NODE_A,
+ case.state.instances,
+ case.state.runners,
+ case.state.tasks,
+ )
assert op == case.expected_op
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 f5a2ac5a..dce20444 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
@@ -27,6 +27,7 @@ from exo.worker.tests.constants import COMMAND_1_ID, INSTANCE_1_ID, MODEL_A_ID
class RunnerSpecDict(TypedDict):
"""Type definition for runner specification dictionaries."""
+
runner_id: RunnerId
node_id: NodeId
device_rank: int
@@ -36,6 +37,7 @@ class RunnerSpecDict(TypedDict):
class MessageDict(TypedDict):
"""Type definition for message dictionaries."""
+
role: Literal["system", "user", "assistant", "developer", "tool", "function"]
content: NotRequired[str | None]
name: NotRequired[str | None]
@@ -46,12 +48,17 @@ class MessageDict(TypedDict):
class TaskSpecDict(TypedDict):
"""Type definition for task specification dictionaries."""
+
task_id: TaskId
- instance_id: NotRequired[InstanceId] # defaults to function parameter if not provided
- command_id: NotRequired[CommandId] # defaults to COMMAND_1_ID if not provided
+ instance_id: NotRequired[
+ InstanceId
+ ] # defaults to function parameter if not provided
+ command_id: NotRequired[CommandId] # defaults to COMMAND_1_ID if not provided
status: NotRequired[TaskStatus] # defaults to TaskStatus.PENDING if not provided
model: NotRequired[str] # defaults to model_id if not provided
- messages: NotRequired[list[MessageDict]] # defaults to [{'role': 'user', 'content': 'Hello, world!'}] if not provided
+ messages: NotRequired[
+ list[MessageDict]
+ ] # defaults to [{'role': 'user', 'content': 'Hello, world!'}] if not provided
@dataclass(slots=True, frozen=True)
@@ -79,10 +86,12 @@ class PlanTestCase:
return self.description.replace(" ", "_")
-def make_shard_metadata(device_rank: int, world_size: int, model_id: ModelId = MODEL_A_ID) -> PipelineShardMetadata:
+def make_shard_metadata(
+ device_rank: int, world_size: int, model_id: ModelId = MODEL_A_ID
+) -> PipelineShardMetadata:
"""Create PipelineShardMetadata with proper layer assignments based on device_rank and world_size."""
total_layers = world_size # For simplicity in tests, total_layers = world_size
-
+
if world_size == 1:
start_layer = 0
end_layer = 1
@@ -92,7 +101,7 @@ def make_shard_metadata(device_rank: int, world_size: int, model_id: ModelId = M
start_layer = device_rank
end_layer = device_rank + 1
n_layers = total_layers
-
+
return PipelineShardMetadata(
device_rank=device_rank,
world_size=world_size,
@@ -112,9 +121,8 @@ def make_downloading_status(node_id: NodeId) -> DownloadingRunnerStatus:
)
)
-def make_model_meta(
- model_id: str
-) -> ModelMetadata:
+
+def make_model_meta(model_id: str) -> ModelMetadata:
model_card: ModelCard
for card in MODEL_CARDS.values():
if card.model_id == model_id:
@@ -126,14 +134,13 @@ def make_model_meta(
storage_size_kilobytes=10**6,
n_layers=16,
)
-
- raise Exception(f'Unknown model_id passed: {model_id}')
+
+ raise Exception(f"Unknown model_id passed: {model_id}")
## Alternatively, if we are ok for this method to be async:
# await _get_model_meta(model_id)
-
def make_instance(
instance_id: InstanceId,
runner_specs: list[tuple[RunnerId, NodeId, int, RunnerStatus]],
@@ -146,11 +153,7 @@ def make_instance(
world_size = len(runner_specs)
for runner_id, node_id, device_rank, _ in runner_specs:
- shard_metadata = make_shard_metadata(
- device_rank,
- world_size,
- model_id
- )
+ shard_metadata = make_shard_metadata(device_rank, world_size, model_id)
runner_to_shard[runner_id] = shard_metadata
node_to_runner[node_id] = runner_id
@@ -167,7 +170,7 @@ def make_instance(
)
# Currently nodes are only ever idle - as if they were running we would be blocking - so we wouldn't be running plan()
- # node_statuses = {node_id: NodeStatus.Idle for _, node_id, _, _ in runner_specs}
+ # node_statuses = {node_id: NodeStatus.Idle for _, node_id, _, _ in runner_specs}
node_statuses: dict[NodeId, NodeStatus] = {}
for _runner_id, node_id, _, status in runner_specs:
if isinstance(status, RunningRunnerStatus):
@@ -178,8 +181,11 @@ def make_instance(
return instance, runner_statuses, node_statuses
+
def make_state(
- runner_specs_per_instance: dict[InstanceId, list[tuple[RunnerId, NodeId, int, RunnerStatus]]],
+ runner_specs_per_instance: dict[
+ InstanceId, list[tuple[RunnerId, NodeId, int, RunnerStatus]]
+ ],
tasks: dict[TaskId, ChatCompletionTask] | None = None,
model_id: ModelId = MODEL_A_ID,
instance_status: InstanceStatus = InstanceStatus.ACTIVE,
@@ -210,6 +216,7 @@ def make_state(
tasks=tasks,
)
+
def make_test_case(
description: str,
runner_specs: list[RunnerSpecDict],
@@ -225,7 +232,7 @@ def make_test_case(
tasks = []
# Convert runner_specs to tuple format for make_instance
specs_tuple = [
- (r['runner_id'], r['node_id'], r['device_rank'], r['status'])
+ (r["runner_id"], r["node_id"], r["device_rank"], r["status"])
for r in runner_specs
]
@@ -234,16 +241,21 @@ def make_test_case(
for t in tasks:
task = ChatCompletionTask(
instance_id=instance_id,
- task_id=t['task_id'],
- command_id=t.get('command_id', command_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=[ChatCompletionMessage(**m) for m in t.get('messages', [{'role': 'user', 'content': 'Hello, world!'}])],
+ model=t.get("model", str(model_id)),
+ messages=[
+ ChatCompletionMessage(**m)
+ for m in t.get(
+ "messages", [{"role": "user", "content": "Hello, world!"}]
+ )
+ ],
),
)
- state_tasks[t['task_id']] = task
+ state_tasks[t["task_id"]] = task
state = make_state(
runner_specs_per_instance={instance_id: specs_tuple},
@@ -255,13 +267,14 @@ def make_test_case(
# Build in_process_runners with downloaded (default True if missing)
in_process_runners = [
InProcessRunner(
- runner_id=r['runner_id'],
+ runner_id=r["runner_id"],
instance_id=instance_id,
model_id=model_id,
- status=r['status'],
- downloaded=r.get('downloaded', True),
- device_rank=r['device_rank'],
- ) for r in runner_specs
+ status=r["status"],
+ downloaded=r.get("downloaded", True),
+ device_rank=r["device_rank"],
+ )
+ for r in runner_specs
]
return PlanTestCase(
@@ -269,4 +282,4 @@ def make_test_case(
state=state,
in_process_runners=in_process_runners,
expected_op=expected_op,
- )
\ No newline at end of file
+ )
diff --git a/src/exo/worker/tests/test_runner_connection.py b/src/exo/worker/tests/test_runner_connection.py
index 80d4c530..196c2401 100644
--- a/src/exo/worker/tests/test_runner_connection.py
+++ b/src/exo/worker/tests/test_runner_connection.py
@@ -29,9 +29,10 @@ from exo.worker.worker import Worker
def user_message() -> str:
return "What is the capital of Japan?"
+
@pytest.mark.skipif(
os.environ.get("DETAILED", "").lower() != "true",
- reason="This test only runs when ENABLE_SPINUP_TIMEOUT_TEST=true environment variable is set"
+ reason="This test only runs when ENABLE_SPINUP_TIMEOUT_TEST=true environment variable is set",
)
async def check_runner_connection(
logger: Logger,
@@ -41,7 +42,7 @@ async def check_runner_connection(
# Track all tasks and workers for cleanup
tasks: list[asyncio.Task[None]] = []
workers: list[Worker] = []
-
+
try:
event_log_manager = EventLogManager(EventLogConfig(), logger)
await event_log_manager.initialize()
@@ -72,49 +73,46 @@ async def check_runner_connection(
task2 = asyncio.create_task(run(worker2, logger))
tasks.append(task2)
- model_id = ModelId('mlx-community/Llama-3.2-1B-Instruct-4bit')
+ model_id = ModelId("mlx-community/Llama-3.2-1B-Instruct-4bit")
shard_assignments = ShardAssignments(
model_id=model_id,
runner_to_shard={
RUNNER_1_ID: pipeline_shard_meta(2, 0),
- RUNNER_2_ID: pipeline_shard_meta(2, 1)
+ RUNNER_2_ID: pipeline_shard_meta(2, 1),
},
- node_to_runner={
- NODE_A: RUNNER_1_ID,
- NODE_B: RUNNER_2_ID
- }
+ node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID},
)
instance = Instance(
instance_id=INSTANCE_1_ID,
instance_type=InstanceStatus.ACTIVE,
shard_assignments=shard_assignments,
- hosts=hosts(2)
+ hosts=hosts(2),
)
await global_events.append_events(
[
- InstanceCreated(
- instance=instance
- ),
+ InstanceCreated(instance=instance),
],
- origin=MASTER_NODE_ID
+ origin=MASTER_NODE_ID,
)
from exo.worker.runner.runner_supervisor import RunnerSupervisor
- async def wait_for_runner_supervisor(worker: Worker, timeout: float = 5.0) -> RunnerSupervisor | None:
+ async def wait_for_runner_supervisor(
+ worker: Worker, timeout: float = 5.0
+ ) -> RunnerSupervisor | None:
end = asyncio.get_event_loop().time() + timeout
while True:
assigned_runners = list(worker.assigned_runners.values())
if assigned_runners:
runner = assigned_runners[0].runner
if isinstance(runner, RunnerSupervisor):
- print('breaking because success')
+ print("breaking because success")
return runner
if isinstance(assigned_runners[0].status, FailedRunnerStatus):
- print('breaking because failed')
+ print("breaking because failed")
return runner
if asyncio.get_event_loop().time() > end:
raise TimeoutError("RunnerSupervisor was not set within timeout")
@@ -129,7 +127,7 @@ async def check_runner_connection(
instance_id=instance.instance_id,
),
],
- origin=MASTER_NODE_ID
+ origin=MASTER_NODE_ID,
)
await asyncio.sleep(0.5)
@@ -139,10 +137,11 @@ async def check_runner_connection(
# Cancel all worker tasks
for task in tasks:
task.cancel()
-
+
# Wait for cancellation to complete
await asyncio.gather(*tasks, return_exceptions=True)
+
# Check Running status
# # not now.
@@ -165,7 +164,7 @@ async def check_runner_connection(
# ) -> None:
# total_runs = 100
# successes = 0
-
+
# for _ in range(total_runs):
# # Create a fresh event loop for each iteration
# loop = asyncio.new_event_loop()
@@ -174,7 +173,7 @@ async def check_runner_connection(
# # Create a fresh event loop for each iteration
# loop = asyncio.new_event_loop()
# asyncio.set_event_loop(loop)
-
+
# try:
# result = loop.run_until_complete(check_runner_connection(
# logger=logger,
@@ -203,16 +202,16 @@ async def check_runner_connection(
# pending = asyncio.all_tasks(loop)
# for task in pending:
# task.cancel()
-
+
# # Run the event loop briefly to allow cancellation to complete
# loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
# # Run the event loop briefly to allow cancellation to complete
# loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
-
+
# # Close the event loop
# loop.close()
# # Close the event loop
# loop.close()
-
+
# print(f"Runner connection successes: {successes} / {total_runs}")
# print(f"Runner connection successes: {successes} / {total_runs}")
diff --git a/src/exo/worker/tests/test_spinup_timeout.py b/src/exo/worker/tests/test_spinup_timeout.py
index b46eb73e..8649fef9 100644
--- a/src/exo/worker/tests/test_spinup_timeout.py
+++ b/src/exo/worker/tests/test_spinup_timeout.py
@@ -20,29 +20,31 @@ from exo.worker.tests.constants import RUNNER_1_ID
# To enable this test, run pytest with: ENABLE_SPINUP_TIMEOUT_TEST=true pytest
+
@pytest.mark.skipif(
os.environ.get("DETAILED", "").lower() != "true",
- reason="This test only runs when ENABLE_SPINUP_TIMEOUT_TEST=true environment variable is set"
+ reason="This test only runs when ENABLE_SPINUP_TIMEOUT_TEST=true environment variable is set",
)
@pytest.mark.asyncio
async def test_runner_up_op_timeout(
- worker_with_assigned_runner: tuple[Worker, Instance],
- chat_completion_task: Callable[[InstanceId, TaskId], Task],
- monkeypatch: pytest.MonkeyPatch
- ):
+ worker_with_assigned_runner: tuple[Worker, Instance],
+ chat_completion_task: Callable[[InstanceId, TaskId], Task],
+ monkeypatch: pytest.MonkeyPatch,
+):
worker, _ = worker_with_assigned_runner
runner_up_op = RunnerUpOp(runner_id=RUNNER_1_ID)
# _execute_runner_up_op should throw a TimeoutError with a short timeout
events: list[Event] = []
- async for event in worker._execute_runner_up_op(runner_up_op, initialize_timeout=0.2): # type: ignore[misc]
+ async for event in worker._execute_runner_up_op(
+ runner_up_op, initialize_timeout=0.2
+ ): # type: ignore[misc]
events.append(event)
assert isinstance(events[-1], RunnerStatusUpdated)
assert isinstance(events[-1].runner_status, FailedRunnerStatus)
assert events[-1].runner_status.error_message is not None
- assert 'timeout' in events[-1].runner_status.error_message.lower()
+ assert "timeout" in events[-1].runner_status.error_message.lower()
del worker.assigned_runners[list(worker.assigned_runners.keys())[0]]
-
diff --git a/src/exo/worker/tests/test_supervisor/test_memory.py b/src/exo/worker/tests/test_supervisor/test_memory.py
index 5eb97b5f..58b3238a 100644
--- a/src/exo/worker/tests/test_supervisor/test_memory.py
+++ b/src/exo/worker/tests/test_supervisor/test_memory.py
@@ -23,9 +23,11 @@ def get_memory_mb(process: Process) -> float:
rss_bytes: int = ps.memory_info().rss # type: ignore[attr-defined]
return rss_bytes / (1024 * 1024)
+
@pytest.fixture
async def model_meta() -> ModelMetadata:
- return await get_model_meta('mlx-community/Llama-3.3-70B-Instruct-4bit')
+ return await get_model_meta("mlx-community/Llama-3.3-70B-Instruct-4bit")
+
@pytest.mark.asyncio
async def test_supervisor_inference_exception(
@@ -45,16 +47,16 @@ async def test_supervisor_inference_exception(
process: Process = supervisor.runner_process
memory = get_memory_mb(process)
- assert memory > 30*100
+ assert memory > 30 * 100
task = chat_completion_task(INSTANCE_1_ID, TASK_1_ID)
- task.task_params.messages[0].content = 'EXO RUNNER MUST FAIL'
+ task.task_params.messages[0].content = "EXO RUNNER MUST FAIL"
with pytest.raises(RunnerError):
async for _ in supervisor.stream_response(task):
pass
await supervisor.astop()
-
+
available_memory_bytes: int = psutil.virtual_memory().available
print(available_memory_bytes // (2**30))
- assert available_memory_bytes > 30 * 2**30
\ No newline at end of file
+ assert available_memory_bytes > 30 * 2**30
diff --git a/src/exo/worker/tests/test_supervisor/test_oom.py b/src/exo/worker/tests/test_supervisor/test_oom.py
index aa2cb6bb..afade315 100644
--- a/src/exo/worker/tests/test_supervisor/test_oom.py
+++ b/src/exo/worker/tests/test_supervisor/test_oom.py
@@ -21,7 +21,9 @@ def user_message():
@pytest.mark.asyncio
-@pytest.mark.skip(reason="Must run `sudo sysctl -w iogpu.wired_limit_mb=` and `sudo sysctl -w iogpu.wired_lwm_mb=` before running this test.")
+@pytest.mark.skip(
+ reason="Must run `sudo sysctl -w iogpu.wired_limit_mb=` and `sudo sysctl -w iogpu.wired_lwm_mb=` before running this test."
+)
async def test_supervisor_catches_oom(
pipeline_shard_meta: Callable[..., PipelineShardMetadata],
hosts: Callable[..., list[Host]],
@@ -38,12 +40,12 @@ async def test_supervisor_catches_oom(
)
task = chat_completion_task(INSTANCE_1_ID, TASK_1_ID)
- task.task_params.messages[0].content = 'EXO RUNNER MUST OOM'
+ task.task_params.messages[0].content = "EXO RUNNER MUST OOM"
with pytest.raises(RunnerError) as exc_info:
async for _ in supervisor.stream_response(task):
pass
-
+
error = exc_info.value
- assert 'memory' in error.error_message.lower()
+ assert "memory" in error.error_message.lower()
await supervisor.astop()
diff --git a/src/exo/worker/tests/test_supervisor/test_supervisor.py b/src/exo/worker/tests/test_supervisor/test_supervisor.py
index 5faf3e57..3452044c 100644
--- a/src/exo/worker/tests/test_supervisor/test_supervisor.py
+++ b/src/exo/worker/tests/test_supervisor/test_supervisor.py
@@ -35,7 +35,7 @@ async def test_supervisor_single_node_response(
model_shard_meta = pipeline_shard_meta(1, 0)
instance_id = InstanceId()
- print(f'{model_shard_meta=}')
+ print(f"{model_shard_meta=}")
supervisor = await RunnerSupervisor.create(
model_shard_meta=model_shard_meta,
@@ -47,7 +47,9 @@ async def test_supervisor_single_node_response(
full_response = ""
stop_reason: FinishReason | None = None
- async for chunk in supervisor.stream_response(task=chat_completion_task(instance_id, TaskId())):
+ async for chunk in supervisor.stream_response(
+ task=chat_completion_task(instance_id, TaskId())
+ ):
if isinstance(chunk, TokenChunk):
full_response += chunk.text
if chunk.finish_reason:
@@ -72,7 +74,7 @@ async def test_supervisor_two_node_response(
):
"""Test that asking for the capital of France returns 'Paris' in the response"""
instance_id = InstanceId()
-
+
async def create_supervisor(shard_idx: int) -> RunnerSupervisor:
supervisor = await RunnerSupervisor.create(
model_shard_meta=pipeline_shard_meta(2, shard_idx),
@@ -80,10 +82,12 @@ async def test_supervisor_two_node_response(
logger=logger,
)
return supervisor
-
+
create_supervisor_0 = asyncio.create_task(create_supervisor(0))
create_supervisor_1 = asyncio.create_task(create_supervisor(1))
- supervisor_0, supervisor_1 = await asyncio.gather(create_supervisor_0, create_supervisor_1)
+ supervisor_0, supervisor_1 = await asyncio.gather(
+ create_supervisor_0, create_supervisor_1
+ )
await asyncio.sleep(0.1)
@@ -93,13 +97,17 @@ async def test_supervisor_two_node_response(
async def collect_response_0():
nonlocal full_response_0
- async for chunk in supervisor_0.stream_response(task=chat_completion_task(instance_id, TaskId())):
+ async for chunk in supervisor_0.stream_response(
+ task=chat_completion_task(instance_id, TaskId())
+ ):
if isinstance(chunk, TokenChunk):
full_response_0 += chunk.text
async def collect_response_1():
nonlocal full_response_1
- async for chunk in supervisor_1.stream_response(task=chat_completion_task(instance_id, TaskId())):
+ async for chunk in supervisor_1.stream_response(
+ task=chat_completion_task(instance_id, TaskId())
+ ):
if isinstance(chunk, TokenChunk):
full_response_1 += chunk.text
@@ -139,11 +147,11 @@ async def test_supervisor_early_stopping(
logger=logger,
)
- task = chat_completion_task(instance_id, TaskId())
+ task = chat_completion_task(instance_id, TaskId())
max_tokens = 50
assert task.task_type == TaskType.CHAT_COMPLETION
- print(f'chat_completion_task.task_params: {task.task_params}')
+ 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/test_supervisor/test_supervisor_sad.py b/src/exo/worker/tests/test_supervisor/test_supervisor_sad.py
index bffae9f5..bfaf1580 100644
--- a/src/exo/worker/tests/test_supervisor/test_supervisor_sad.py
+++ b/src/exo/worker/tests/test_supervisor/test_supervisor_sad.py
@@ -29,6 +29,7 @@ async def test_supervisor_instantiation_exception(
logger=logger,
)
+
@pytest.mark.asyncio
async def test_supervisor_instantiation_timeout(
pipeline_shard_meta: Callable[..., PipelineShardMetadata],
@@ -37,7 +38,7 @@ async def test_supervisor_instantiation_timeout(
):
"""Test that asking for the capital of France returns 'Paris' in the response"""
model_shard_meta = pipeline_shard_meta(1, 0)
- model_shard_meta.should_timeout = 10 # timeout after 10s
+ model_shard_meta.should_timeout = 10 # timeout after 10s
with pytest.raises(asyncio.TimeoutError):
_ = await RunnerSupervisor.create(
@@ -47,7 +48,6 @@ async def test_supervisor_instantiation_timeout(
)
-
@pytest.mark.asyncio
async def test_supervisor_inference_exception(
pipeline_shard_meta: Callable[..., PipelineShardMetadata],
@@ -65,11 +65,12 @@ async def test_supervisor_inference_exception(
)
task = chat_completion_task(INSTANCE_1_ID, TASK_1_ID)
- task.task_params.messages[0].content = 'EXO RUNNER MUST FAIL'
+ task.task_params.messages[0].content = "EXO RUNNER MUST FAIL"
with pytest.raises(RunnerError):
async for _ in supervisor.stream_response(task):
pass
+
@pytest.mark.asyncio
async def test_supervisor_inference_timeout(
pipeline_shard_meta: Callable[..., PipelineShardMetadata],
@@ -87,9 +88,9 @@ async def test_supervisor_inference_timeout(
)
task = chat_completion_task(INSTANCE_1_ID, TASK_1_ID)
- task.task_params.messages[0].content = 'EXO RUNNER MUST TIMEOUT'
+ task.task_params.messages[0].content = "EXO RUNNER MUST TIMEOUT"
with pytest.raises(asyncio.TimeoutError):
async for _ in supervisor.stream_response(task):
pass
- await asyncio.sleep(0.1)
\ No newline at end of file
+ await asyncio.sleep(0.1)
diff --git a/src/exo/worker/utils/macmon/__init__.py b/src/exo/worker/utils/macmon/__init__.py
index ad950d89..bf4bda58 100644
--- a/src/exo/worker/utils/macmon/__init__.py
+++ b/src/exo/worker/utils/macmon/__init__.py
@@ -1,3 +1,3 @@
from .macmon import MacMonError, get_metrics, get_metrics_async
-__all__ = ['get_metrics', 'get_metrics_async', 'MacMonError']
\ No newline at end of file
+__all__ = ["get_metrics", "get_metrics_async", "MacMonError"]
diff --git a/src/exo/worker/utils/macmon/macmon.py b/src/exo/worker/utils/macmon/macmon.py
index 8814fbd9..1d823c2a 100644
--- a/src/exo/worker/utils/macmon/macmon.py
+++ b/src/exo/worker/utils/macmon/macmon.py
@@ -1,7 +1,7 @@
import asyncio
import platform
-import subprocess
import shutil
+import subprocess
from typing import Optional, Tuple
from pydantic import BaseModel, ConfigDict, ValidationError
@@ -27,11 +27,10 @@ def _get_binary_path() -> str:
):
raise MacMonError("MacMon only supports macOS with Apple Silicon (ARM) chips")
-
path = shutil.which("macmon")
if path is None:
- raise MacMonError(f"MacMon not found in PATH")
+ raise MacMonError("MacMon not found in PATH")
return path
diff --git a/src/exo/worker/utils/profile.py b/src/exo/worker/utils/profile.py
index d1763eb3..be4a17ea 100644
--- a/src/exo/worker/utils/profile.py
+++ b/src/exo/worker/utils/profile.py
@@ -67,7 +67,7 @@ async def start_polling_node_metrics(
# Run heavy FLOPs profiling only if enough time has elapsed
- override_memory_env = os.getenv('OVERRIDE_MEMORY')
+ override_memory_env = os.getenv("OVERRIDE_MEMORY")
override_memory: int | None = (
int(override_memory_env) * 2**30 if override_memory_env else None
)
@@ -80,7 +80,9 @@ async def start_polling_node_metrics(
network_interfaces=network_interfaces,
memory=MemoryPerformanceProfile(
ram_total=total_mem,
- ram_available=override_memory if override_memory else total_mem - used_mem,
+ ram_available=override_memory
+ if override_memory
+ else total_mem - used_mem,
swap_total=metrics.memory.swap_total
if metrics.memory is not None
and metrics.memory.swap_total is not None
@@ -94,12 +96,25 @@ async def start_polling_node_metrics(
),
system=SystemPerformanceProfile(
flops_fp16=0,
- gpu_usage=metrics.gpu_usage[1] if metrics.gpu_usage is not None else 0,
- temp=metrics.temp.gpu_temp_avg if metrics.temp is not None and metrics.temp.gpu_temp_avg is not None else 0,
- sys_power=metrics.sys_power if metrics.sys_power is not None else 0,
- pcpu_usage=metrics.pcpu_usage[1] if metrics.pcpu_usage is not None else 0,
- ecpu_usage=metrics.ecpu_usage[1] if metrics.ecpu_usage is not None else 0,
- ane_power=metrics.ane_power if metrics.ane_power is not None else 0,
+ gpu_usage=metrics.gpu_usage[1]
+ if metrics.gpu_usage is not None
+ else 0,
+ temp=metrics.temp.gpu_temp_avg
+ if metrics.temp is not None
+ and metrics.temp.gpu_temp_avg is not None
+ else 0,
+ sys_power=metrics.sys_power
+ if metrics.sys_power is not None
+ else 0,
+ pcpu_usage=metrics.pcpu_usage[1]
+ if metrics.pcpu_usage is not None
+ else 0,
+ ecpu_usage=metrics.ecpu_usage[1]
+ if metrics.ecpu_usage is not None
+ else 0,
+ ane_power=metrics.ane_power
+ if metrics.ane_power is not None
+ else 0,
),
)
)
diff --git a/src/exo/worker/utils/system_info.py b/src/exo/worker/utils/system_info.py
index b5e2a0c8..1aa69325 100644
--- a/src/exo/worker/utils/system_info.py
+++ b/src/exo/worker/utils/system_info.py
@@ -21,7 +21,7 @@ async def get_mac_friendly_name_async() -> str | None:
e.g., "John's MacBook Pro"
Returns the name as a string, or None if an error occurs or not on macOS.
"""
- if sys.platform != 'darwin': # 'darwin' is the platform name for macOS
+ if sys.platform != "darwin": # 'darwin' is the platform name for macOS
print("This function is designed for macOS only.")
return None
@@ -30,9 +30,11 @@ async def get_mac_friendly_name_async() -> str | None:
# stdout=asyncio.subprocess.PIPE captures standard output.
# stderr=asyncio.subprocess.PIPE captures standard error.
process = await asyncio.create_subprocess_exec(
- 'scutil', '--get', 'ComputerName',
+ "scutil",
+ "--get",
+ "ComputerName",
stdout=asyncio.subprocess.PIPE,
- stderr=asyncio.subprocess.PIPE
+ stderr=asyncio.subprocess.PIPE,
)
# process.communicate() reads all data from stdout and stderr
@@ -52,8 +54,12 @@ async def get_mac_friendly_name_async() -> str | None:
return None
else:
# If there was an error, print the stderr output
- error_message = stderr_data.decode().strip() if stderr_data else "Unknown error"
- print(f"Error executing scutil (return code {process.returncode}): {error_message}")
+ error_message = (
+ stderr_data.decode().strip() if stderr_data else "Unknown error"
+ )
+ print(
+ f"Error executing scutil (return code {process.returncode}): {error_message}"
+ )
return None
except FileNotFoundError:
@@ -64,6 +70,7 @@ async def get_mac_friendly_name_async() -> str | None:
print(f"An unexpected error occurred: {e}")
return None
+
async def get_network_interface_info_async() -> List[NetworkInterfaceInfo]:
"""
Retrieves detailed network interface information on macOS.
@@ -71,7 +78,7 @@ async def get_network_interface_info_async() -> List[NetworkInterfaceInfo]:
to determine interface names, IP addresses, and types (ethernet, wifi, vpn, other).
Returns a list of NetworkInterfaceInfo objects.
"""
- if sys.platform != 'darwin':
+ if sys.platform != "darwin":
return []
interfaces_info: List[NetworkInterfaceInfo] = []
@@ -83,25 +90,37 @@ async def get_network_interface_info_async() -> List[NetworkInterfaceInfo]:
process = await asyncio.create_subprocess_exec(
*command_parts,
stdout=asyncio.subprocess.PIPE,
- stderr=asyncio.subprocess.PIPE
+ stderr=asyncio.subprocess.PIPE,
)
stdout_data, stderr_data = await process.communicate()
if process.returncode == 0:
# Use 'utf-8' and replace errors for robustness
- return stdout_data.decode('utf-8', errors='replace').strip()
+ return stdout_data.decode("utf-8", errors="replace").strip()
else:
- error_message = stderr_data.decode('utf-8', errors='replace').strip() if stderr_data else "Unknown error"
- print(f"Error executing {' '.join(command_parts)} (code {process.returncode}): {error_message}")
+ error_message = (
+ stderr_data.decode("utf-8", errors="replace").strip()
+ if stderr_data
+ else "Unknown error"
+ )
+ print(
+ f"Error executing {' '.join(command_parts)} (code {process.returncode}): {error_message}"
+ )
return None
except FileNotFoundError:
- print(f"Error: Command '{command_parts[0]}' not found. Ensure it's in PATH.")
+ print(
+ f"Error: Command '{command_parts[0]}' not found. Ensure it's in PATH."
+ )
return None
except Exception as e:
- print(f"An unexpected error occurred running {' '.join(command_parts)}: {e}")
+ print(
+ f"An unexpected error occurred running {' '.join(command_parts)}: {e}"
+ )
return None
# 1. Get hardware port types from networksetup
- networksetup_output = await _run_cmd_async(['networksetup', '-listallhardwareports'])
+ networksetup_output = await _run_cmd_async(
+ ["networksetup", "-listallhardwareports"]
+ )
if networksetup_output:
current_hardware_port_type_raw: Optional[str] = None
for line in networksetup_output.splitlines():
@@ -112,46 +131,49 @@ async def get_network_interface_info_async() -> List[NetworkInterfaceInfo]:
device_name = line_stripped.split(":", 1)[1].strip()
if device_name and device_name != "N/A":
if "Thunderbolt" in current_hardware_port_type_raw:
- device_to_type_map[device_name] = 'thunderbolt'
- elif "Wi-Fi" in current_hardware_port_type_raw or "AirPort" in current_hardware_port_type_raw:
- device_to_type_map[device_name] = 'wifi'
- elif "Ethernet" in current_hardware_port_type_raw or \
- "LAN" in current_hardware_port_type_raw:
- device_to_type_map[device_name] = 'ethernet'
- current_hardware_port_type_raw = None # Reset for the next block
+ device_to_type_map[device_name] = "thunderbolt"
+ elif (
+ "Wi-Fi" in current_hardware_port_type_raw
+ or "AirPort" in current_hardware_port_type_raw
+ ):
+ device_to_type_map[device_name] = "wifi"
+ elif (
+ "Ethernet" in current_hardware_port_type_raw
+ or "LAN" in current_hardware_port_type_raw
+ ):
+ device_to_type_map[device_name] = "ethernet"
+ current_hardware_port_type_raw = None # Reset for the next block
# 2. Get interface names and IP addresses from ifconfig
- ifconfig_output = await _run_cmd_async(['ifconfig'])
+ ifconfig_output = await _run_cmd_async(["ifconfig"])
if ifconfig_output:
current_if_name: Optional[str] = None
# Regex for interface name (e.g., en0:, utun0:, tailscale0.)
- interface_header_pattern = re.compile(r'^([a-zA-Z0-9\._-]+):')
+ interface_header_pattern = re.compile(r"^([a-zA-Z0-9\._-]+):")
# Regex for IPv4 address (inet)
- inet_pattern = re.compile(r'^\s+inet\s+(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})')
+ inet_pattern = re.compile(r"^\s+inet\s+(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})")
# Regex for IPv6 address (inet6)
- inet6_pattern = re.compile(r'^\s+inet6\s+([0-9a-fA-F:]+(?:%[a-zA-Z0-9._-]+)?)')
+ inet6_pattern = re.compile(r"^\s+inet6\s+([0-9a-fA-F:]+(?:%[a-zA-Z0-9._-]+)?)")
def _add_interface_entry(if_name: str, ip_addr: str):
_if_type = device_to_type_map.get(if_name)
- if not _if_type: # Infer type if not found via networksetup
+ if not _if_type: # Infer type if not found via networksetup
if if_name.startswith(("utun", "wg", "ppp")) or "tailscale" in if_name:
- _if_type = 'vpn'
+ _if_type = "vpn"
elif if_name.startswith("bridge"):
- _if_type = 'virtual' # For non-Thunderbolt bridges (e.g., Docker)
+ _if_type = "virtual" # For non-Thunderbolt bridges (e.g., Docker)
else:
- _if_type = 'other'
-
- interfaces_info.append(NetworkInterfaceInfo(
- name=if_name,
- ip_address=ip_addr,
- type=_if_type
- ))
+ _if_type = "other"
+
+ interfaces_info.append(
+ NetworkInterfaceInfo(name=if_name, ip_address=ip_addr, type=_if_type)
+ )
for line in ifconfig_output.splitlines():
header_match = interface_header_pattern.match(line)
if header_match:
potential_if_name = header_match.group(1)
- if potential_if_name == "lo0": # Skip loopback interface
+ if potential_if_name == "lo0": # Skip loopback interface
current_if_name = None
else:
current_if_name = potential_if_name
@@ -161,7 +183,9 @@ async def get_network_interface_info_async() -> List[NetworkInterfaceInfo]:
inet_m = inet_pattern.match(line)
if inet_m:
ipv4_address = inet_m.group(1)
- _add_interface_entry(current_if_name, ipv4_address) # Add all IPv4, including APIPA
+ _add_interface_entry(
+ current_if_name, ipv4_address
+ ) # Add all IPv4, including APIPA
continue
inet6_m = inet6_pattern.match(line)
@@ -169,9 +193,10 @@ async def get_network_interface_info_async() -> List[NetworkInterfaceInfo]:
ipv6_address = inet6_m.group(1)
# No specific filtering for IPv6 link-local (e.g., fe80::) for now.
_add_interface_entry(current_if_name, ipv6_address)
-
+
return interfaces_info
+
async def get_mac_system_info_async() -> SystemInfo:
"""Get Mac system information using system_profiler."""
model_id_val = "Unknown Model"
@@ -181,42 +206,61 @@ async def get_mac_system_info_async() -> SystemInfo:
try:
process = await asyncio.create_subprocess_exec(
- "system_profiler", "SPHardwareDataType",
+ "system_profiler",
+ "SPHardwareDataType",
stdout=asyncio.subprocess.PIPE,
- stderr=asyncio.subprocess.PIPE
+ stderr=asyncio.subprocess.PIPE,
)
stdout_data, stderr_data = await process.communicate()
if process.returncode == 0:
if stdout_data:
output = stdout_data.decode().strip()
- model_line = next((line for line in output.split("\n") if "Model Name" in line), None)
- model_id_val = model_line.split(": ")[1] if model_line else "Unknown Model"
+ model_line = next(
+ (line for line in output.split("\n") if "Model Name" in line), None
+ )
+ model_id_val = (
+ model_line.split(": ")[1] if model_line else "Unknown Model"
+ )
- chip_line = next((line for line in output.split("\n") if "Chip" in line), None)
+ chip_line = next(
+ (line for line in output.split("\n") if "Chip" in line), None
+ )
chip_id_val = chip_line.split(": ")[1] if chip_line else "Unknown Chip"
- memory_line = next((line for line in output.split("\n") if "Memory" in line), None)
- memory_str = memory_line.split(": ")[1] if memory_line else "0 GB" # Default to "0 GB"
+ memory_line = next(
+ (line for line in output.split("\n") if "Memory" in line), None
+ )
+ memory_str = (
+ memory_line.split(": ")[1] if memory_line else "0 GB"
+ ) # Default to "0 GB"
memory_units = memory_str.split()
if len(memory_units) == 2:
try:
memory_value_int = int(memory_units[0])
if memory_units[1] == "GB":
- memory_val = memory_value_int * 1024 # Assuming MB
+ memory_val = memory_value_int * 1024 # Assuming MB
elif memory_units[1] == "MB":
- memory_val = memory_value_int
- else: # TB? Unlikely for typical memory, handle gracefully
- memory_val = memory_value_int # Store as is, let consumer decide unit or log
+ memory_val = memory_value_int
+ else: # TB? Unlikely for typical memory, handle gracefully
+ memory_val = memory_value_int # Store as is, let consumer decide unit or log
print(f"Warning: Unknown memory unit {memory_units[1]}")
except ValueError:
- print(f"Warning: Could not parse memory value {memory_units[0]}")
+ print(
+ f"Warning: Could not parse memory value {memory_units[0]}"
+ )
memory_val = 0
else:
- print("system_profiler command succeeded but produced no output for hardware.")
+ print(
+ "system_profiler command succeeded but produced no output for hardware."
+ )
else:
- error_message = stderr_data.decode().strip() if stderr_data else "Unknown error"
- print(f"Error executing system_profiler (return code {process.returncode}): {error_message}")
+ error_message = (
+ stderr_data.decode().strip() if stderr_data else "Unknown error"
+ )
+ print(
+ f"Error executing system_profiler (return code {process.returncode}): {error_message}"
+ )
except Exception as e:
print(f"Error getting Mac hardware info: {e}")
@@ -227,10 +271,9 @@ async def get_mac_system_info_async() -> SystemInfo:
print(f"Error getting Mac network interface info: {e}")
network_interfaces_info_list = []
-
return SystemInfo(
model_id=model_id_val,
chip_id=chip_id_val,
memory=memory_val,
- network_interfaces=network_interfaces_info_list
+ network_interfaces=network_interfaces_info_list,
)
diff --git a/src/exo/worker/worker.py b/src/exo/worker/worker.py
index 25dfd36b..0c66dc76 100644
--- a/src/exo/worker/worker.py
+++ b/src/exo/worker/worker.py
@@ -60,7 +60,9 @@ class Worker:
self.node_id: NodeId = node_id
self.state: State = State()
self.shard_downloader: ShardDownloader = shard_downloader
- self.worker_events: AsyncSQLiteEventStorage | None = worker_events # worker_events is None in some tests.
+ self.worker_events: AsyncSQLiteEventStorage | None = (
+ worker_events # worker_events is None in some tests.
+ )
self.global_events: AsyncSQLiteEventStorage | None = global_events
self.logger: logging.Logger = logger
@@ -100,11 +102,16 @@ class Worker:
self, assigned_runner: AssignedRunner
) -> AsyncGenerator[Event, None]:
"""Handles the case where the shard is already downloaded."""
- async for event in self._update_runner_status_to_completed_then_inactive(assigned_runner):
+ async for event in self._update_runner_status_to_completed_then_inactive(
+ assigned_runner
+ ):
yield event
async def _handle_shard_download_process(
- self, assigned_runner: AssignedRunner, op: AssignRunnerOp, initial_progress: RepoDownloadProgress
+ self,
+ assigned_runner: AssignedRunner,
+ op: AssignRunnerOp,
+ initial_progress: RepoDownloadProgress,
) -> AsyncGenerator[Event, None]:
"""Manages the shard download process with progress tracking."""
# Set initial ongoing status
@@ -113,8 +120,8 @@ class Worker:
node_id=self.node_id,
download_progress=DownloadProgressData(
total_bytes=initial_progress.total_bytes,
- downloaded_bytes=initial_progress.downloaded_bytes
- )
+ downloaded_bytes=initial_progress.downloaded_bytes,
+ ),
)
)
yield assigned_runner.status_update_event()
@@ -122,31 +129,45 @@ class Worker:
# Set up download progress tracking
download_progress_queue: asyncio.Queue[RepoDownloadProgress] = asyncio.Queue()
- def download_progress_callback(shard: ShardMetadata, progress: RepoDownloadProgress) -> None:
+ def download_progress_callback(
+ shard: ShardMetadata, progress: RepoDownloadProgress
+ ) -> None:
download_progress_queue.put_nowait(progress)
self.shard_downloader.on_progress(download_progress_callback)
- download_task = asyncio.create_task(self.shard_downloader.ensure_shard(op.shard_metadata))
+ download_task = asyncio.create_task(
+ self.shard_downloader.ensure_shard(op.shard_metadata)
+ )
try:
- async for event in self._monitor_download_progress(assigned_runner, download_progress_queue):
+ async for event in self._monitor_download_progress(
+ assigned_runner, download_progress_queue
+ ):
yield event
finally:
if not download_task.done():
download_task.cancel()
async def _monitor_download_progress(
- self, assigned_runner: AssignedRunner, download_progress_queue: asyncio.Queue[RepoDownloadProgress]
+ self,
+ assigned_runner: AssignedRunner,
+ download_progress_queue: asyncio.Queue[RepoDownloadProgress],
) -> AsyncGenerator[Event, None]:
"""Monitors download progress and yields status updates."""
last_progress_time = 0.0
throttle_interval_secs = 1.0
while True:
- progress: RepoDownloadProgress = await asyncio.wait_for(download_progress_queue.get(), timeout=15)
+ progress: RepoDownloadProgress = await asyncio.wait_for(
+ download_progress_queue.get(), timeout=15
+ )
if progress.status == "complete":
- async for event in self._update_runner_status_to_completed_then_inactive(assigned_runner):
+ async for (
+ event
+ ) in self._update_runner_status_to_completed_then_inactive(
+ assigned_runner
+ ):
yield event
break
elif progress.status == "in_progress":
@@ -157,7 +178,7 @@ class Worker:
download_progress=DownloadProgressData(
total_bytes=progress.total_bytes,
downloaded_bytes=progress.downloaded_bytes,
- )
+ ),
)
)
yield assigned_runner.status_update_event()
@@ -171,13 +192,19 @@ class Worker:
This op assigns the runner, and moves from Downloading -> Inactive (ready to spin) state.
"""
assigned_runner = self._create_assigned_runner(op)
- initial_progress = await self.shard_downloader.get_shard_download_status_for_shard(op.shard_metadata)
+ initial_progress = (
+ await self.shard_downloader.get_shard_download_status_for_shard(
+ op.shard_metadata
+ )
+ )
if initial_progress.status == "complete":
async for event in self._handle_already_downloaded_shard(assigned_runner):
yield event
else:
- async for event in self._handle_shard_download_process(assigned_runner, op, initial_progress):
+ async for event in self._handle_shard_download_process(
+ assigned_runner, op, initial_progress
+ ):
yield event
async def _execute_unassign_op(
@@ -207,7 +234,7 @@ class Worker:
model_shard_meta=assigned_runner.shard_metadata,
hosts=assigned_runner.hosts,
logger=self.logger,
- initialize_timeout=initialize_timeout
+ initialize_timeout=initialize_timeout,
)
if assigned_runner.runner.healthy:
@@ -216,17 +243,21 @@ class Worker:
# Log detailed reasons why the runner is not healthy
runner = assigned_runner.runner
health_issues: list[str] = []
-
+
if runner.runner_process.returncode is not None:
- health_issues.append(f"runner_process.returncode is {runner.runner_process.returncode}")
+ health_issues.append(
+ f"runner_process.returncode is {runner.runner_process.returncode}"
+ )
if runner.runner_process.stdin is None:
health_issues.append("runner_process.stdin is None")
elif runner.runner_process.stdin.is_closing():
health_issues.append("runner_process.stdin is closing")
if runner.runner_process.stdout is None:
health_issues.append("runner_process.stdout is None")
-
- self.logger.warning(f"Runner status is not healthy: {', '.join(health_issues)}")
+
+ self.logger.warning(
+ f"Runner status is not healthy: {', '.join(health_issues)}"
+ )
assigned_runner.status = FailedRunnerStatus()
yield self.assigned_runners[op.runner_id].status_update_event()
@@ -247,29 +278,28 @@ class Worker:
async def _execute_runner_failed_op(
self, op: RunnerFailedOp
) -> AsyncGenerator[Event, None]:
- '''
+ """
We detected that this runner has failed. So we'll put it into 'failed' state now, triggering the rest of the instance to spin down.
- '''
+ """
assigned_runner = self.assigned_runners[op.runner_id]
if isinstance(assigned_runner.runner, RunnerSupervisor):
- await assigned_runner.runner.astop() # astop the runner to ensure it clears out of memory.
+ await (
+ assigned_runner.runner.astop()
+ ) # astop the runner to ensure it clears out of memory.
assigned_runner.status = FailedRunnerStatus()
yield self.assigned_runners[op.runner_id].status_update_event()
-
- async def _execute_task_op(
- self, op: ExecuteTaskOp
- ) -> AsyncGenerator[Event, None]:
- '''
+ async def _execute_task_op(self, op: ExecuteTaskOp) -> AsyncGenerator[Event, None]:
+ """
This is the entry point for a chat completion starting.
While there is only one execute function, it will get called in different ways for runner 0 and runner [1, 2, 3, ...].
Runners [1, 2, 3, ...] will run this method when a task is in 'pending' state.
Runner 0 will run this method when a task is in 'running' state.
TODO: How do we handle the logic of ensuring that n-1 nodes have started their execution before allowing the 0'th runner to start?
This is still a little unclear to me.
- '''
+ """
assigned_runner = self.assigned_runners[op.runner_id]
async def inner_execute(queue: asyncio.Queue[Event]) -> None:
@@ -279,36 +309,41 @@ class Worker:
await queue.put(assigned_runner.status_update_event())
if assigned_runner.shard_metadata.device_rank == 0:
- await queue.put(TaskStateUpdated(
- task_id=op.task.task_id,
- task_status=TaskStatus.RUNNING,
- ))
+ await queue.put(
+ TaskStateUpdated(
+ task_id=op.task.task_id,
+ task_status=TaskStatus.RUNNING,
+ )
+ )
assert assigned_runner.runner is not None
assert assigned_runner.runner.healthy
async for chunk in assigned_runner.runner.stream_response(
- task=op.task,
- request_started_callback=partial(running_callback, queue)):
+ task=op.task, request_started_callback=partial(running_callback, queue)
+ ):
if assigned_runner.shard_metadata.device_rank == 0:
- await queue.put(ChunkGenerated(
- # todo: at some point we will no longer have a bijection between task_id and row_id.
- # So we probably want to store a mapping between these two in our Worker object.
- command_id=chunk.command_id,
- chunk=chunk
- ))
+ await queue.put(
+ ChunkGenerated(
+ # todo: at some point we will no longer have a bijection between task_id and row_id.
+ # So we probably want to store a mapping between these two in our Worker object.
+ command_id=chunk.command_id,
+ chunk=chunk,
+ )
+ )
if assigned_runner.shard_metadata.device_rank == 0:
- await queue.put(TaskStateUpdated(
- task_id=op.task.task_id,
- task_status=TaskStatus.COMPLETE,
- ))
+ await queue.put(
+ TaskStateUpdated(
+ task_id=op.task.task_id,
+ task_status=TaskStatus.COMPLETE,
+ )
+ )
# After a successful inference:
assigned_runner.status = LoadedRunnerStatus()
await queue.put(assigned_runner.status_update_event())
-
queue: Queue[Event] = asyncio.Queue()
task = asyncio.create_task(inner_execute(queue))
@@ -340,8 +375,9 @@ class Worker:
try:
await asyncio.wait_for(task, timeout=5)
except asyncio.TimeoutError:
- self.logger.warning("Timed out waiting for task cleanup after inference execution.")
-
+ self.logger.warning(
+ "Timed out waiting for task cleanup after inference execution."
+ )
## Operation Planner
@@ -364,31 +400,26 @@ class Worker:
async for event in event_generator:
yield event
-
- async def fail_runner(self, e: Exception, runner_id: RunnerId) -> AsyncGenerator[Event]:
+ async def fail_runner(
+ self, e: Exception, runner_id: RunnerId
+ ) -> AsyncGenerator[Event]:
if runner_id in self.assigned_runners:
assigned_runner = self.assigned_runners[runner_id]
assigned_runner.runner = None
assigned_runner.status = FailedRunnerStatus(error_message=str(e))
- assigned_runner.failures.append(
- (
- time.time(),
- e
- )
- )
+ assigned_runner.failures.append((time.time(), e))
# Reset failure count back to 0 when succesful
if len(assigned_runner.failures) >= 3:
# Too many retries. We will emit a DeleteInstance
- yield InstanceDeleted(
- instance_id=assigned_runner.instance_id
- )
+ yield InstanceDeleted(instance_id=assigned_runner.instance_id)
yield assigned_runner.status_update_event()
-
- async def fail_task(self, e: Exception, runner_id: RunnerId, task_id: TaskId) -> AsyncGenerator[Event]:
+ async def fail_task(
+ self, e: Exception, runner_id: RunnerId, task_id: TaskId
+ ) -> AsyncGenerator[Event]:
if runner_id in self.assigned_runners:
yield TaskStateUpdated(
task_id=task_id,
@@ -396,15 +427,12 @@ class Worker:
)
yield TaskFailed(
- task_id=task_id,
- error_type=str(type(e)),
- error_message=str(e)
+ task_id=task_id, error_type=str(type(e)), error_message=str(e)
)
async for event in self.fail_runner(e, runner_id):
yield event
-
async def event_publisher(self, event: Event) -> None:
assert self.worker_events is not None
await self.worker_events.append_events([event], self.node_id)
← be6f5ae7 feat: build system and homebrew compatibility
·
back to Exo
·
add EXO logo to dashboard 5bfc99b4 →