← back to Exo
Refactor tasks / commands / api
7ac23ce96b11b617f459b027684175fdbdbba397 · 2025-07-23 15:52:29 +0100 · Seth Howes
Files touched
M master/api.pyM master/main.pyA networking/target/rust-analyzer/metadata/sysroot/Cargo.lockD shared/event_loops/commands.pyD shared/event_loops/main.pyM shared/tests/test_sqlite_connector.pyM shared/types/api.pyM shared/types/events/_events.pyM shared/types/events/chunks.pyM shared/types/events/commands.pyD shared/types/request.pyM shared/types/state.pyM shared/types/tasks.pyM worker/main.pyM worker/runner/runner_supervisor.pyM worker/tests/conftest.py
Diff
commit 7ac23ce96b11b617f459b027684175fdbdbba397
Author: Seth Howes <71157822+sethhowes@users.noreply.github.com>
Date: Wed Jul 23 15:52:29 2025 +0100
Refactor tasks / commands / api
---
master/api.py | 47 +-
master/main.py | 16 +-
.../rust-analyzer/metadata/sysroot/Cargo.lock | 503 +++++++++++++++++++++
shared/event_loops/commands.py | 28 --
shared/event_loops/main.py | 120 -----
shared/tests/test_sqlite_connector.py | 12 +-
shared/types/api.py | 66 ++-
shared/types/events/_events.py | 4 +-
shared/types/events/chunks.py | 4 +-
shared/types/events/commands.py | 52 ++-
shared/types/request.py | 12 -
shared/types/state.py | 9 +-
shared/types/tasks.py | 23 +-
worker/main.py | 7 +-
worker/runner/runner_supervisor.py | 4 +-
worker/tests/conftest.py | 18 +-
16 files changed, 658 insertions(+), 267 deletions(-)
diff --git a/master/api.py b/master/api.py
index f07e81f5..ec697140 100644
--- a/master/api.py
+++ b/master/api.py
@@ -2,38 +2,30 @@ import asyncio
import time
from asyncio.queues import Queue
from collections.abc import AsyncGenerator
-from typing import List, Optional, Sequence, final
+from typing import Sequence, final
import uvicorn
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
-from pydantic import BaseModel
from shared.db.sqlite.connector import AsyncSQLiteEventStorage
+from shared.types.api import (
+ ChatCompletionMessage,
+ ChatCompletionResponse,
+ StreamingChoiceResponse,
+)
from shared.types.events import ChunkGenerated, Event
from shared.types.events.chunks import TokenChunk
+from shared.types.events.commands import (
+ ChatCompletionCommand,
+ Command,
+ CommandId,
+ CommandTypes,
+)
from shared.types.events.components import EventFromEventLog
-from shared.types.request import APIRequest, RequestId
from shared.types.tasks import ChatCompletionTaskParams
-class Message(BaseModel):
- role: str
- content: str
-
-class StreamingChoiceResponse(BaseModel):
- index: int
- delta: Message
- finish_reason: Optional[str] = None
-
-
-class ChatCompletionResponse(BaseModel):
- id: str
- object: str = "chat.completion"
- created: int
- model: str
- choices: List[StreamingChoiceResponse]
-
def chunk_to_response(chunk: TokenChunk) -> ChatCompletionResponse:
return ChatCompletionResponse(
id='abc',
@@ -42,7 +34,7 @@ def chunk_to_response(chunk: TokenChunk) -> ChatCompletionResponse:
choices=[
StreamingChoiceResponse(
index=0,
- delta=Message(
+ delta=ChatCompletionMessage(
role='assistant',
content=chunk.text
),
@@ -54,7 +46,7 @@ def chunk_to_response(chunk: TokenChunk) -> ChatCompletionResponse:
@final
class API:
- def __init__(self, command_queue: Queue[APIRequest], global_events: AsyncSQLiteEventStorage) -> None:
+ def __init__(self, command_queue: Queue[Command], global_events: AsyncSQLiteEventStorage) -> None:
self._app = FastAPI()
self._setup_routes()
@@ -106,10 +98,11 @@ class API:
# At the moment, we just create the task in the API.
# In the future, a `Request` will be created here and they will be bundled into `Task` objects by the master.
- request_id=RequestId()
+ command_id=CommandId()
- request = APIRequest(
- request_id=request_id,
+ request = ChatCompletionCommand(
+ command_id=command_id,
+ command_type=CommandTypes.CHAT_COMPLETION,
request_params=payload,
)
await self.command_queue.put(request)
@@ -124,7 +117,7 @@ class API:
for wrapped_event in events:
event = wrapped_event.event
- if isinstance(event, ChunkGenerated) and event.request_id == request_id:
+ if isinstance(event, ChunkGenerated) and event.command_id == command_id:
assert isinstance(event.chunk, TokenChunk)
chunk_response: ChatCompletionResponse = chunk_to_response(event.chunk)
print(chunk_response)
@@ -146,7 +139,7 @@ class API:
def start_fastapi_server(
- command_queue: Queue[APIRequest],
+ command_queue: Queue[Command],
global_events: AsyncSQLiteEventStorage,
host: str = "0.0.0.0",
port: int = 8000,
diff --git a/master/main.py b/master/main.py
index 3e99f808..6cb646aa 100644
--- a/master/main.py
+++ b/master/main.py
@@ -10,11 +10,11 @@ from shared.db.sqlite.event_log_manager import EventLogManager
from shared.types.common import NodeId
from shared.types.events import ChunkGenerated
from shared.types.events.chunks import TokenChunk
-from shared.types.request import APIRequest, RequestId
+from shared.types.events.commands import Command, CommandId
## TODO: Hook this up properly
-async def fake_tokens_task(events_log: AsyncSQLiteEventStorage, request_id: RequestId):
+async def fake_tokens_task(events_log: AsyncSQLiteEventStorage, command_id: CommandId):
model_id = "testmodelabc"
for i in range(10):
@@ -22,9 +22,9 @@ async def fake_tokens_task(events_log: AsyncSQLiteEventStorage, request_id: Requ
# Create the event with proper types and consistent IDs
chunk_event = ChunkGenerated(
- request_id=request_id,
+ command_id=command_id,
chunk=TokenChunk(
- request_id=request_id, # Use the same task_id
+ command_id=command_id, # Use the same task_id
idx=i,
model=model_id, # Use the same model_id
text=f'text{i}',
@@ -42,9 +42,9 @@ async def fake_tokens_task(events_log: AsyncSQLiteEventStorage, request_id: Requ
# Create the event with proper types and consistent IDs
chunk_event = ChunkGenerated(
- request_id=request_id,
+ command_id=command_id,
chunk=TokenChunk(
- request_id=request_id, # Use the same task_id
+ command_id=command_id, # Use the same task_id
idx=11,
model=model_id, # Use the same model_id
text=f'text{11}',
@@ -68,7 +68,7 @@ async def main():
await event_log_manager.initialize()
global_events: AsyncSQLiteEventStorage = event_log_manager.global_events
- command_queue: Queue[APIRequest] = asyncio.Queue()
+ command_queue: Queue[Command] = asyncio.Queue()
api_thread = threading.Thread(
target=start_fastapi_server,
@@ -88,7 +88,7 @@ async def main():
print(command)
- await fake_tokens_task(global_events, request_id=command.request_id)
+ await fake_tokens_task(global_events, command_id=command.command_id)
await asyncio.sleep(0.01)
diff --git a/networking/target/rust-analyzer/metadata/sysroot/Cargo.lock b/networking/target/rust-analyzer/metadata/sysroot/Cargo.lock
new file mode 100644
index 00000000..97996d5f
--- /dev/null
+++ b/networking/target/rust-analyzer/metadata/sysroot/Cargo.lock
@@ -0,0 +1,503 @@
+# This file is automatically @generated by Cargo.
+# It is not intended for manual editing.
+version = 4
+
+[[package]]
+name = "addr2line"
+version = "0.22.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6e4503c46a5c0c7844e948c9a4d6acd9f50cccb4de1c48eb9e291ea17470c678"
+dependencies = [
+ "compiler_builtins",
+ "gimli 0.29.0",
+ "rustc-std-workspace-alloc",
+ "rustc-std-workspace-core",
+]
+
+[[package]]
+name = "adler"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
+dependencies = [
+ "compiler_builtins",
+ "rustc-std-workspace-core",
+]
+
+[[package]]
+name = "alloc"
+version = "0.0.0"
+dependencies = [
+ "compiler_builtins",
+ "core",
+ "rand",
+ "rand_xorshift",
+]
+
+[[package]]
+name = "allocator-api2"
+version = "0.2.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5c6cb57a04249c6480766f7f7cef5467412af1490f8d1e243141daddada3264f"
+
+[[package]]
+name = "cc"
+version = "1.1.22"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9540e661f81799159abee814118cc139a2004b3a3aa3ea37724a1b66530b90e0"
+dependencies = [
+ "shlex",
+]
+
+[[package]]
+name = "cfg-if"
+version = "1.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
+dependencies = [
+ "compiler_builtins",
+ "rustc-std-workspace-core",
+]
+
+[[package]]
+name = "compiler_builtins"
+version = "0.1.138"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "53f0ea7fff95b51f84371588f06062557e96bbe363d2b36218ddb806f3ca8611"
+dependencies = [
+ "cc",
+ "rustc-std-workspace-core",
+]
+
+[[package]]
+name = "core"
+version = "0.0.0"
+dependencies = [
+ "rand",
+ "rand_xorshift",
+]
+
+[[package]]
+name = "dlmalloc"
+version = "0.2.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d9b5e0d321d61de16390ed273b647ce51605b575916d3c25e6ddf27a1e140035"
+dependencies = [
+ "cfg-if",
+ "compiler_builtins",
+ "libc",
+ "rustc-std-workspace-core",
+ "windows-sys",
+]
+
+[[package]]
+name = "fortanix-sgx-abi"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "57cafc2274c10fab234f176b25903ce17e690fca7597090d50880e047a0389c5"
+dependencies = [
+ "compiler_builtins",
+ "rustc-std-workspace-core",
+]
+
+[[package]]
+name = "getopts"
+version = "0.2.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "14dbbfd5c71d70241ecf9e6f13737f7b5ce823821063188d7e46c41d371eebd5"
+dependencies = [
+ "rustc-std-workspace-core",
+ "rustc-std-workspace-std",
+ "unicode-width",
+]
+
+[[package]]
+name = "gimli"
+version = "0.29.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "40ecd4077b5ae9fd2e9e169b102c6c330d0605168eb0e8bf79952b256dbefffd"
+dependencies = [
+ "compiler_builtins",
+ "rustc-std-workspace-alloc",
+ "rustc-std-workspace-core",
+]
+
+[[package]]
+name = "gimli"
+version = "0.31.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f"
+dependencies = [
+ "compiler_builtins",
+ "rustc-std-workspace-alloc",
+ "rustc-std-workspace-core",
+]
+
+[[package]]
+name = "hashbrown"
+version = "0.15.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1e087f84d4f86bf4b218b927129862374b72199ae7d8657835f1e89000eea4fb"
+dependencies = [
+ "allocator-api2",
+ "compiler_builtins",
+ "rustc-std-workspace-alloc",
+ "rustc-std-workspace-core",
+]
+
+[[package]]
+name = "hermit-abi"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fbf6a919d6cf397374f7dfeeea91d974c7c0a7221d0d0f4f20d859d329e53fcc"
+dependencies = [
+ "compiler_builtins",
+ "rustc-std-workspace-alloc",
+ "rustc-std-workspace-core",
+]
+
+[[package]]
+name = "libc"
+version = "0.2.162"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "18d287de67fe55fd7e1581fe933d965a5a9477b38e949cfa9f8574ef01506398"
+dependencies = [
+ "rustc-std-workspace-core",
+]
+
+[[package]]
+name = "memchr"
+version = "2.7.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3"
+dependencies = [
+ "compiler_builtins",
+ "rustc-std-workspace-core",
+]
+
+[[package]]
+name = "miniz_oxide"
+version = "0.7.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b8a240ddb74feaf34a79a7add65a741f3167852fba007066dcac1ca548d89c08"
+dependencies = [
+ "adler",
+ "compiler_builtins",
+ "rustc-std-workspace-alloc",
+ "rustc-std-workspace-core",
+]
+
+[[package]]
+name = "object"
+version = "0.36.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "aedf0a2d09c573ed1d8d85b30c119153926a2b36dce0ab28322c09a117a4683e"
+dependencies = [
+ "compiler_builtins",
+ "memchr",
+ "rustc-std-workspace-alloc",
+ "rustc-std-workspace-core",
+]
+
+[[package]]
+name = "panic_abort"
+version = "0.0.0"
+dependencies = [
+ "alloc",
+ "cfg-if",
+ "compiler_builtins",
+ "core",
+ "libc",
+]
+
+[[package]]
+name = "panic_unwind"
+version = "0.0.0"
+dependencies = [
+ "alloc",
+ "cfg-if",
+ "compiler_builtins",
+ "core",
+ "libc",
+ "unwind",
+]
+
+[[package]]
+name = "proc_macro"
+version = "0.0.0"
+dependencies = [
+ "core",
+ "std",
+]
+
+[[package]]
+name = "profiler_builtins"
+version = "0.0.0"
+dependencies = [
+ "cc",
+ "compiler_builtins",
+ "core",
+]
+
+[[package]]
+name = "r-efi"
+version = "4.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e9e935efc5854715dfc0a4c9ef18dc69dee0ec3bf9cc3ab740db831c0fdd86a3"
+dependencies = [
+ "compiler_builtins",
+ "rustc-std-workspace-core",
+]
+
+[[package]]
+name = "r-efi-alloc"
+version = "1.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "31d6f09fe2b6ad044bc3d2c34ce4979796581afd2f1ebc185837e02421e02fd7"
+dependencies = [
+ "compiler_builtins",
+ "r-efi",
+ "rustc-std-workspace-core",
+]
+
+[[package]]
+name = "rand"
+version = "0.8.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404"
+dependencies = [
+ "rand_core",
+]
+
+[[package]]
+name = "rand_core"
+version = "0.6.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
+
+[[package]]
+name = "rand_xorshift"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d25bf25ec5ae4a3f1b92f929810509a2f53d7dca2f50b794ff57e3face536c8f"
+dependencies = [
+ "rand_core",
+]
+
+[[package]]
+name = "rustc-demangle"
+version = "0.1.24"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "719b953e2095829ee67db738b3bfa9fa368c94900df327b3f07fe6e794d2fe1f"
+dependencies = [
+ "compiler_builtins",
+ "rustc-std-workspace-core",
+]
+
+[[package]]
+name = "rustc-std-workspace-alloc"
+version = "1.99.0"
+dependencies = [
+ "alloc",
+]
+
+[[package]]
+name = "rustc-std-workspace-core"
+version = "1.99.0"
+dependencies = [
+ "core",
+]
+
+[[package]]
+name = "rustc-std-workspace-std"
+version = "1.99.0"
+dependencies = [
+ "std",
+]
+
+[[package]]
+name = "shlex"
+version = "1.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64"
+
+[[package]]
+name = "std"
+version = "0.0.0"
+dependencies = [
+ "addr2line",
+ "alloc",
+ "cfg-if",
+ "compiler_builtins",
+ "core",
+ "dlmalloc",
+ "fortanix-sgx-abi",
+ "hashbrown",
+ "hermit-abi",
+ "libc",
+ "miniz_oxide",
+ "object",
+ "panic_abort",
+ "panic_unwind",
+ "r-efi",
+ "r-efi-alloc",
+ "rand",
+ "rand_xorshift",
+ "rustc-demangle",
+ "std_detect",
+ "unwind",
+ "wasi",
+ "windows-targets 0.0.0",
+]
+
+[[package]]
+name = "std_detect"
+version = "0.1.5"
+dependencies = [
+ "cfg-if",
+ "compiler_builtins",
+ "libc",
+ "rustc-std-workspace-alloc",
+ "rustc-std-workspace-core",
+]
+
+[[package]]
+name = "sysroot"
+version = "0.0.0"
+dependencies = [
+ "proc_macro",
+ "profiler_builtins",
+ "std",
+ "test",
+]
+
+[[package]]
+name = "test"
+version = "0.0.0"
+dependencies = [
+ "core",
+ "getopts",
+ "libc",
+ "std",
+]
+
+[[package]]
+name = "unicode-width"
+version = "0.1.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7dd6e30e90baa6f72411720665d41d89b9a3d039dc45b8faea1ddd07f617f6af"
+dependencies = [
+ "compiler_builtins",
+ "rustc-std-workspace-core",
+ "rustc-std-workspace-std",
+]
+
+[[package]]
+name = "unwind"
+version = "0.0.0"
+dependencies = [
+ "cfg-if",
+ "compiler_builtins",
+ "core",
+ "libc",
+ "unwinding",
+]
+
+[[package]]
+name = "unwinding"
+version = "0.2.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "637d511437df708cee34bdec7ba2f1548d256b7acf3ff20e0a1c559f9bf3a987"
+dependencies = [
+ "compiler_builtins",
+ "gimli 0.31.1",
+ "rustc-std-workspace-core",
+]
+
+[[package]]
+name = "wasi"
+version = "0.11.0+wasi-snapshot-preview1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423"
+dependencies = [
+ "compiler_builtins",
+ "rustc-std-workspace-alloc",
+ "rustc-std-workspace-core",
+]
+
+[[package]]
+name = "windows-sys"
+version = "0.59.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b"
+dependencies = [
+ "windows-targets 0.52.6",
+]
+
+[[package]]
+name = "windows-targets"
+version = "0.0.0"
+
+[[package]]
+name = "windows-targets"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973"
+dependencies = [
+ "windows_aarch64_gnullvm",
+ "windows_aarch64_msvc",
+ "windows_i686_gnu",
+ "windows_i686_gnullvm",
+ "windows_i686_msvc",
+ "windows_x86_64_gnu",
+ "windows_x86_64_gnullvm",
+ "windows_x86_64_msvc",
+]
+
+[[package]]
+name = "windows_aarch64_gnullvm"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3"
+
+[[package]]
+name = "windows_aarch64_msvc"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469"
+
+[[package]]
+name = "windows_i686_gnu"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b"
+
+[[package]]
+name = "windows_i686_gnullvm"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66"
+
+[[package]]
+name = "windows_i686_msvc"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66"
+
+[[package]]
+name = "windows_x86_64_gnu"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78"
+
+[[package]]
+name = "windows_x86_64_gnullvm"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d"
+
+[[package]]
+name = "windows_x86_64_msvc"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec"
diff --git a/shared/event_loops/commands.py b/shared/event_loops/commands.py
deleted file mode 100644
index ac79b3b8..00000000
--- a/shared/event_loops/commands.py
+++ /dev/null
@@ -1,28 +0,0 @@
-from typing import Annotated, Literal
-
-from pydantic import BaseModel, Field, TypeAdapter
-
-from shared.types.common import NewUUID
-
-
-class ExternalCommandId(NewUUID):
- pass
-
-
-class BaseExternalCommand[T: str](BaseModel):
- command_id: ExternalCommandId
- command_type: T
-
-
-class ChatCompletionNonStreamingCommand(
- BaseExternalCommand[Literal["chat_completion_non_streaming"]]
-):
- command_type: Literal["chat_completion_non_streaming"] = (
- "chat_completion_non_streaming"
- )
-
-
-ExternalCommand = Annotated[
- ChatCompletionNonStreamingCommand, Field(discriminator="command_type")
-]
-ExternalCommandParser: TypeAdapter[ExternalCommand] = TypeAdapter(ExternalCommand)
diff --git a/shared/event_loops/main.py b/shared/event_loops/main.py
deleted file mode 100644
index 582745e6..00000000
--- a/shared/event_loops/main.py
+++ /dev/null
@@ -1,120 +0,0 @@
-from asyncio import Lock, Task
-from asyncio import Queue as AsyncQueue
-from collections.abc import MutableMapping
-from logging import Logger
-from typing import Any, Hashable, Mapping, Protocol, Sequence
-
-from fastapi.responses import Response, StreamingResponse
-
-from shared.event_loops.commands import ExternalCommand
-from shared.types.events import Event
-from shared.types.events.components import Apply, EventFromEventLog
-from shared.types.state import State
-
-
-class ExhaustiveMapping[K: Hashable, V](MutableMapping[K, V]):
- __slots__ = ("_store",)
-
- required_keys: frozenset[K] = frozenset()
-
- def __init__(self, data: Mapping[K, V]):
- missing = self.required_keys - data.keys()
- extra = data.keys() - self.required_keys
- if missing or extra:
- raise ValueError(f"missing={missing!r}, extra={extra!r}")
- self._store: dict[K, V] = dict(data)
-
- def __getitem__(self, k: K) -> V:
- return self._store[k]
-
- def __setitem__(self, k: K, v: V) -> None:
- self._store[k] = v
-
- def __delitem__(self, k: K) -> None:
- del self._store[k]
-
- def __iter__(self):
- return iter(self._store)
-
- def __len__(self) -> int:
- return len(self._store)
-
-
-def apply_events(
- state: State, apply_fn: Apply, events: Sequence[EventFromEventLog[Event]]
-) -> State:
- sorted_events = sorted(events, key=lambda event: event.idx_in_log)
- state = state.model_copy()
- for wrapped_event in sorted_events:
- if wrapped_event.idx_in_log <= state.last_event_applied_idx:
- continue
- state.last_event_applied_idx = wrapped_event.idx_in_log
- state = apply_fn(state, wrapped_event.event)
- return state
-
-
-class NodeCommandLoopProtocol(Protocol):
- _command_runner: Task[Any] | None = None
- _command_queue: AsyncQueue[ExternalCommand]
- _response_queue: AsyncQueue[Response | StreamingResponse]
- _logger: Logger
-
- @property
- def is_command_runner_running(self) -> bool:
- return self._command_runner is not None and not self._command_runner.done()
-
- async def start_command_runner(self) -> None: ...
- async def stop_command_runner(self) -> None: ...
- async def push_command(self, command: ExternalCommand) -> None: ...
- async def pop_response(self) -> Response | StreamingResponse: ...
- async def _handle_command(self, command: ExternalCommand) -> None: ...
-
-
-class NodeEventGetterProtocol(Protocol):
- _event_fetcher: Task[Any] | None = None
- _event_queue: AsyncQueue[EventFromEventLog[Event]]
- _logger: Logger
-
- @property
- async def is_event_fetcher_running(self) -> bool:
- return self._event_fetcher is not None and not self._event_fetcher.done()
-
- async def start_event_fetcher(self) -> None: ...
- async def stop_event_fetcher(self) -> None: ...
-
-
-class NodeStateStorageProtocol(Protocol):
- _state: State
- _state_lock: Lock
- _logger: Logger
-
- async def _read_state(
- self,
- ) -> State: ...
-
-
-class NodeStateManagerProtocol(
- NodeEventGetterProtocol, NodeStateStorageProtocol
-):
- _state_manager: Task[Any] | None = None
- _logger: Logger
-
- @property
- async def is_state_manager_running(self) -> bool:
- is_task_running = (
- self._state_manager is not None and not self._state_manager.done()
- )
- return (
- is_task_running
- and await self.is_event_fetcher_running
- and await self.is_state_manager_running
- )
-
- async def start_state_manager(self) -> None: ...
- async def stop_state_manager(self) -> None: ...
- async def _apply_queued_events(self) -> None: ...
-
-
-class NodeEventLoopProtocol(
- NodeCommandLoopProtocol, NodeStateManagerProtocol
-): ...
diff --git a/shared/tests/test_sqlite_connector.py b/shared/tests/test_sqlite_connector.py
index 50fef7ad..80736bfd 100644
--- a/shared/tests/test_sqlite_connector.py
+++ b/shared/tests/test_sqlite_connector.py
@@ -16,7 +16,7 @@ from shared.types.events import (
_EventType,
)
from shared.types.events.chunks import ChunkType, TokenChunk
-from shared.types.request import RequestId
+from shared.types.events.commands import CommandId
# Type ignore comment for all protected member access in this test file
# pyright: reportPrivateUsage=false
@@ -439,19 +439,19 @@ class TestAsyncSQLiteEventStorage:
await storage.start()
# Create a ChunkGenerated event with nested TokenChunk
- request_id = RequestId(uuid=uuid4())
+ command_id = CommandId(uuid=uuid4())
token_chunk = TokenChunk(
text="Hello, world!",
token_id=42,
finish_reason="stop",
chunk_type=ChunkType.token,
- request_id=request_id,
+ command_id=command_id,
idx=0,
model="test-model"
)
chunk_generated_event = ChunkGenerated(
- request_id=request_id,
+ command_id=command_id,
chunk=token_chunk
)
@@ -473,13 +473,13 @@ class TestAsyncSQLiteEventStorage:
retrieved_event = retrieved_event_wrapper.event
assert isinstance(retrieved_event, ChunkGenerated)
assert retrieved_event.event_type == _EventType.ChunkGenerated
- assert retrieved_event.request_id == request_id
+ assert retrieved_event.command_id == command_id
# Verify the nested chunk was deserialized correctly
retrieved_chunk = retrieved_event.chunk
assert isinstance(retrieved_chunk, TokenChunk)
assert retrieved_chunk.chunk_type == ChunkType.token
- assert retrieved_chunk.request_id == request_id
+ assert retrieved_chunk.command_id == command_id
assert retrieved_chunk.idx == 0
assert retrieved_chunk.model == "test-model"
diff --git a/shared/types/api.py b/shared/types/api.py
index 37f1a74e..28adb93d 100644
--- a/shared/types/api.py
+++ b/shared/types/api.py
@@ -2,6 +2,8 @@ from typing import Any, Literal
from pydantic import BaseModel
+from shared.openai_compat import FinishReason
+
class ChatCompletionMessage(BaseModel):
role: Literal["system", "user", "assistant", "developer", "tool", "function"]
@@ -12,6 +14,68 @@ class ChatCompletionMessage(BaseModel):
function_call: dict[str, Any] | None = None
+class TopLogprobItem(BaseModel):
+ token: str
+ logprob: float
+ bytes: list[int] | None = None
+
+
+class LogprobsContentItem(BaseModel):
+ token: str
+ logprob: float
+ bytes: list[int] | None = None
+ top_logprobs: list[TopLogprobItem]
+
+
+class Logprobs(BaseModel):
+ content: list[LogprobsContentItem] | None = None
+
+
+class StreamingChoiceResponse(BaseModel):
+ index: int
+ delta: ChatCompletionMessage
+ logprobs: Logprobs | None = None
+ finish_reason: FinishReason | None = None
+
+
+class ChatCompletionChoice(BaseModel):
+ index: int
+ message: ChatCompletionMessage
+ logprobs: Logprobs | None = None
+ finish_reason: FinishReason | None = None
+
+
+class PromptTokensDetails(BaseModel):
+ cached_tokens: int = 0
+ audio_tokens: int = 0
+
+
+class CompletionTokensDetails(BaseModel):
+ reasoning_tokens: int = 0
+ audio_tokens: int = 0
+ accepted_prediction_tokens: int = 0
+ rejected_prediction_tokens: int = 0
+
+
+class Usage(BaseModel):
+ prompt_tokens: int
+ completion_tokens: int
+ total_tokens: int
+ prompt_tokens_details: PromptTokensDetails | None = None
+ completion_tokens_details: CompletionTokensDetails | None = None
+
+
+
+class ChatCompletionResponse(BaseModel):
+ id: str
+ object: Literal["chat.completion"] = "chat.completion"
+ created: int
+ model: str
+ choices: list[ChatCompletionChoice | StreamingChoiceResponse]
+ usage: Usage | None = None
+ service_tier: str | None = None
+
+
class ChatCompletionTaskParams(BaseModel):
model: str
frequency_penalty: float | None = None
@@ -31,4 +95,4 @@ class ChatCompletionTaskParams(BaseModel):
tools: list[dict[str, Any]] | None = None
tool_choice: str | dict[str, Any] | None = None
parallel_tool_calls: bool | None = None
- user: str | None = None
\ No newline at end of file
+ user: str | None = None
diff --git a/shared/types/events/_events.py b/shared/types/events/_events.py
index 0c3a80f7..64cefe50 100644
--- a/shared/types/events/_events.py
+++ b/shared/types/events/_events.py
@@ -3,7 +3,7 @@ from typing import Literal
from shared.topology import Connection, ConnectionProfile, Node, NodePerformanceProfile
from shared.types.common import NodeId
from shared.types.events.chunks import GenerationChunk
-from shared.types.request import RequestId
+from shared.types.events.commands import CommandId
from shared.types.tasks import Task, TaskId, TaskStatus
from shared.types.worker.common import InstanceId, NodeStatus
from shared.types.worker.instances import InstanceParams, TypeOfInstance
@@ -101,7 +101,7 @@ class WorkerDisconnected(_BaseEvent[_EventType.WorkerDisconnected]):
class ChunkGenerated(_BaseEvent[_EventType.ChunkGenerated]):
event_type: Literal[_EventType.ChunkGenerated] = _EventType.ChunkGenerated
- request_id: RequestId
+ command_id: CommandId
chunk: GenerationChunk
diff --git a/shared/types/events/chunks.py b/shared/types/events/chunks.py
index 9504496a..81d2bfae 100644
--- a/shared/types/events/chunks.py
+++ b/shared/types/events/chunks.py
@@ -4,8 +4,8 @@ from typing import Annotated, Literal
from pydantic import BaseModel, Field, TypeAdapter
from shared.openai_compat import FinishReason
+from shared.types.events.commands import CommandId
from shared.types.models import ModelId
-from shared.types.request import RequestId
class ChunkType(str, Enum):
@@ -15,7 +15,7 @@ class ChunkType(str, Enum):
class BaseChunk[ChunkTypeT: ChunkType](BaseModel):
chunk_type: ChunkTypeT
- request_id: RequestId
+ command_id: CommandId
idx: int
model: ModelId
diff --git a/shared/types/events/commands.py b/shared/types/events/commands.py
index 6651d823..fe645869 100644
--- a/shared/types/events/commands.py
+++ b/shared/types/events/commands.py
@@ -1,19 +1,12 @@
from enum import Enum
-from typing import (
- TYPE_CHECKING,
- Callable,
- Sequence,
-)
+from typing import Annotated, Callable, Sequence
-if TYPE_CHECKING:
- pass
-
-from pydantic import BaseModel
+from pydantic import BaseModel, Field, TypeAdapter
+from shared.types.api import ChatCompletionTaskParams
from shared.types.common import NewUUID
-from shared.types.state import State
-
-from . import Event
+from shared.types.events import Event
+from shared.types.state import InstanceId, State
class CommandId(NewUUID):
@@ -21,19 +14,36 @@ class CommandId(NewUUID):
class CommandTypes(str, Enum):
- Create = "Create"
- Update = "Update"
- Delete = "Delete"
+ CHAT_COMPLETION = "CHAT_COMPLETION"
+ CREATE_INSTANCE = "CREATE_INSTANCE"
+ DELETE_INSTANCE = "DELETE_INSTANCE"
-class Command[
- CommandType: CommandTypes,
-](BaseModel):
- command_type: CommandType
+class _BaseCommand[T: CommandTypes](BaseModel):
command_id: CommandId
+ command_type: T
+
+
+class ChatCompletionCommand(_BaseCommand[CommandTypes.CHAT_COMPLETION]):
+ request_params: ChatCompletionTaskParams
+
+
+class CreateInstanceCommand(_BaseCommand[CommandTypes.CREATE_INSTANCE]):
+ model_id: str
+
+
+class DeleteInstanceCommand(_BaseCommand[CommandTypes.DELETE_INSTANCE]):
+ instance_id: InstanceId
+
+
+Command = Annotated[
+ ChatCompletionCommand, Field(discriminator="command_type")
+]
+
+CommandParser: TypeAdapter[Command] = TypeAdapter(Command)
-type Decide[CommandTypeT: CommandTypes] = Callable[
- [State, Command[CommandTypeT]],
+type Decide = Callable[
+ [State, Command],
Sequence[Event],
]
diff --git a/shared/types/request.py b/shared/types/request.py
deleted file mode 100644
index a9a267a8..00000000
--- a/shared/types/request.py
+++ /dev/null
@@ -1,12 +0,0 @@
-from pydantic import BaseModel
-
-from shared.types.api import ChatCompletionTaskParams
-from shared.types.common import NewUUID
-
-
-class RequestId(NewUUID):
- pass
-
-class APIRequest(BaseModel):
- request_id: RequestId
- request_params: ChatCompletionTaskParams
\ No newline at end of file
diff --git a/shared/types/state.py b/shared/types/state.py
index 5851034c..0129d925 100644
--- a/shared/types/state.py
+++ b/shared/types/state.py
@@ -1,21 +1,17 @@
from collections.abc import Mapping, Sequence
from enum import Enum
-from typing import List
from pydantic import BaseModel, ConfigDict, Field
from shared.topology import Topology
from shared.types.common import NodeId
from shared.types.profiling import NodePerformanceProfile
-from shared.types.tasks import Task, TaskId, TaskSagaEntry
+from shared.types.tasks import Task, TaskId
from shared.types.worker.common import InstanceId, NodeStatus
from shared.types.worker.instances import BaseInstance
from shared.types.worker.runners import RunnerId, RunnerStatus
-class ExternalCommand(BaseModel): ...
-
-
class CachePolicy(str, Enum):
KEEP_ALL = "KEEP_ALL"
@@ -26,11 +22,8 @@ class State(BaseModel):
instances: Mapping[InstanceId, BaseInstance] = {}
runners: Mapping[RunnerId, RunnerStatus] = {}
tasks: Mapping[TaskId, Task] = {}
- task_sagas: Mapping[TaskId, Sequence[TaskSagaEntry]] = {}
node_profiles: Mapping[NodeId, NodePerformanceProfile] = {}
topology: Topology = Topology()
history: Sequence[Topology] = []
- task_inbox: List[Task] = Field(default_factory=list)
- task_outbox: List[Task] = Field(default_factory=list)
cache_policy: CachePolicy = CachePolicy.KEEP_ALL
last_event_applied_idx: int = Field(default=0, ge=0)
diff --git a/shared/types/tasks.py b/shared/types/tasks.py
index 2d865f57..011f084a 100644
--- a/shared/types/tasks.py
+++ b/shared/types/tasks.py
@@ -1,6 +1,7 @@
from enum import Enum
+from typing import Annotated
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
from shared.types.api import ChatCompletionTaskParams
from shared.types.common import NewUUID
@@ -10,25 +11,23 @@ from shared.types.worker.common import InstanceId
class TaskId(NewUUID):
pass
+
class TaskType(str, Enum):
CHAT_COMPLETION = "CHAT_COMPLETION"
+
class TaskStatus(str, Enum):
- Pending = "Pending"
- Running = "Running"
- Complete = "Complete"
- Failed = "Failed"
+ PENDING = "PENDING"
+ RUNNING = "RUNNING"
+ COMPLETE = "COMPLETE"
+ FAILED = "FAILED"
-class Task(BaseModel):
- task_id: TaskId
+class ChatCompletionTask(BaseModel):
task_type: TaskType
+ task_id: TaskId
instance_id: InstanceId
task_status: TaskStatus
task_params: ChatCompletionTaskParams
-
-
-class TaskSagaEntry(BaseModel):
- task_id: TaskId
- instance_id: InstanceId
+Task = Annotated[ChatCompletionTask, Field(discriminator="task_type")]
diff --git a/worker/main.py b/worker/main.py
index 5c73512f..b1274a70 100644
--- a/worker/main.py
+++ b/worker/main.py
@@ -243,7 +243,7 @@ class Worker:
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.
- request_id=chunk.request_id,
+ command_id=chunk.command_id,
chunk=chunk
))
@@ -338,12 +338,9 @@ class Worker:
async def _loop(self):
while True:
state_copy = self.state.model_copy(deep=False)
- state_copy.task_inbox = []
- state_copy.task_outbox = []
-
op: RunnerOp | None = self.plan(state_copy)
- # Run the op, synchronously blocking for now.
+ # run the op, synchronously blocking for now
if op is not None:
async for event in self._execute_op(op):
print(event)
diff --git a/worker/runner/runner_supervisor.py b/worker/runner/runner_supervisor.py
index b889be4b..1f60f1d9 100644
--- a/worker/runner/runner_supervisor.py
+++ b/worker/runner/runner_supervisor.py
@@ -6,7 +6,7 @@ from types import CoroutineType
from typing import Any, Callable
from shared.types.events.chunks import GenerationChunk, TokenChunk
-from shared.types.request import RequestId
+from shared.types.events.commands import CommandId
from shared.types.tasks import ChatCompletionTaskParams, Task
from shared.types.worker.commands_runner import (
ChatTaskMessage,
@@ -181,7 +181,7 @@ class RunnerSupervisor:
text=text, token=token, finish_reason=finish_reason
):
yield TokenChunk(
- request_id=RequestId(uuid=task.task_id.uuid),
+ command_id=CommandId(uuid=task.task_id.uuid),
idx=token,
model=self.model_shard_meta.model_meta.model_id,
text=text,
diff --git a/worker/tests/conftest.py b/worker/tests/conftest.py
index b101a853..7e4f003c 100644
--- a/worker/tests/conftest.py
+++ b/worker/tests/conftest.py
@@ -11,7 +11,7 @@ from shared.types.common import NodeId
from shared.types.models import ModelId, ModelMetadata
from shared.types.state import State
from shared.types.tasks import (
- Task,
+ ChatCompletionTask,
TaskId,
TaskStatus,
TaskType,
@@ -105,21 +105,13 @@ def completion_create_params(user_message: str) -> ChatCompletionTaskParams:
)
@pytest.fixture
-def chat_completion_task(completion_create_params: ChatCompletionTaskParams) -> Task:
- """Creates a ChatCompletionTask directly for serdes testing"""
- return Task(task_id=TaskId(), instance_id=InstanceId(), task_type=TaskType.CHAT_COMPLETION, task_status=TaskStatus.Pending, task_params=completion_create_params)
-
-@pytest.fixture
-def chat_task(
- completion_create_params: ChatCompletionTaskParams,
-) -> Task:
- """Creates the final Task object"""
- return Task(
+def chat_completion_task(completion_create_params: ChatCompletionTaskParams) -> ChatCompletionTask:
+ return ChatCompletionTask(
task_id=TaskId(),
instance_id=InstanceId(),
task_type=TaskType.CHAT_COMPLETION,
- task_status=TaskStatus.Pending,
- task_params=completion_create_params,
+ task_status=TaskStatus.PENDING,
+ task_params=completion_create_params
)
@pytest.fixture
← 81060b70 Made basedpyright work with Jetbrains environment
·
back to Exo
·
Fix tests 7a452c33 →