From 7ac23ce96b11b617f459b027684175fdbdbba397 Mon Sep 17 00:00:00 2001 From: Seth Howes <71157822+sethhowes@users.noreply.github.com> Date: Wed, 23 Jul 2025 15:52:29 +0100 Subject: [PATCH] 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(-) create mode 100644 networking/target/rust-analyzer/metadata/sysroot/Cargo.lock delete mode 100644 shared/event_loops/commands.py delete mode 100644 shared/event_loops/main.py delete mode 100644 shared/types/request.py 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 -type Decide[CommandTypeT: CommandTypes] = Callable[ - [State, Command[CommandTypeT]], +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 = 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