From 07598a3af1f6ce0649ea49de7a25dc698b10434f Mon Sep 17 00:00:00 2001 From: Evan Quiney Date: Thu, 19 Mar 2026 15:57:50 +0000 Subject: [PATCH] teeny refactor (#1753) api.py keeps growing. it's not tied to the master, so should have it's own top level folder. ## testing ci, pytest --- src/exo/api/__init__.py | 0 src/exo/{master => api}/adapters/__init__.py | 0 .../adapters/chat_completions.py | 2 +- src/exo/{master => api}/adapters/claude.py | 16 ++-- src/exo/{master => api}/adapters/ollama.py | 16 ++-- src/exo/{master => api}/adapters/responses.py | 18 ++-- src/exo/{master/api.py => api/main.py} | 96 +++++++++---------- .../tests/test_api_error_handling.py | 3 +- .../tests/test_cancel_command.py | 2 +- .../{master => api}/tests/test_claude_api.py | 4 +- .../tests/test_claude_tool_use.py | 4 +- .../tests/test_openai_responses_api.py | 4 +- src/exo/api/types/__init__.py | 57 +++++++++++ src/exo/{shared => api}/types/api.py | 0 src/exo/{shared => api}/types/claude_api.py | 0 src/exo/{shared => api}/types/ollama_api.py | 0 .../{shared => api}/types/openai_responses.py | 0 src/exo/main.py | 2 +- src/exo/master/main.py | 2 +- src/exo/shared/types/chunks.py | 8 +- src/exo/shared/types/commands.py | 4 +- src/exo/shared/types/tasks.py | 2 +- .../shared/types/worker/runner_response.py | 2 +- .../event_log.py => utils/disk_event_log.py} | 0 src/exo/utils/power_sampler.py | 2 +- .../{master => utils}/tests/test_event_log.py | 2 +- src/exo/utils/tests/test_power_sampler.py | 2 +- .../worker/engines/image/distributed_model.py | 2 +- src/exo/worker/engines/image/generate.py | 2 +- src/exo/worker/engines/mlx/dsml_encoding.py | 2 +- .../engines/mlx/generator/batch_generate.py | 2 +- .../worker/engines/mlx/generator/generate.py | 2 +- src/exo/worker/main.py | 2 +- src/exo/worker/runner/image_models/runner.py | 2 +- .../llm_inference/model_output_parsers.py | 2 +- .../runner/llm_inference/tool_parsers.py | 2 +- .../test_runner/test_parse_gpt_oss.py | 2 +- 37 files changed, 163 insertions(+), 105 deletions(-) create mode 100644 src/exo/api/__init__.py rename src/exo/{master => api}/adapters/__init__.py (100%) rename src/exo/{master => api}/adapters/chat_completions.py (99%) rename src/exo/{master => api}/adapters/claude.py (99%) rename src/exo/{master => api}/adapters/ollama.py (99%) rename src/exo/{master => api}/adapters/responses.py (99%) rename src/exo/{master/api.py => api/main.py} (99%) rename src/exo/{master => api}/tests/test_api_error_handling.py (97%) rename src/exo/{master => api}/tests/test_cancel_command.py (98%) rename src/exo/{master => api}/tests/test_claude_api.py (98%) rename src/exo/{master => api}/tests/test_claude_tool_use.py (99%) rename src/exo/{master => api}/tests/test_openai_responses_api.py (96%) create mode 100644 src/exo/api/types/__init__.py rename src/exo/{shared => api}/types/api.py (100%) rename src/exo/{shared => api}/types/claude_api.py (100%) rename src/exo/{shared => api}/types/ollama_api.py (100%) rename src/exo/{shared => api}/types/openai_responses.py (100%) rename src/exo/{master/event_log.py => utils/disk_event_log.py} (100%) rename src/exo/{master => utils}/tests/test_event_log.py (98%) diff --git a/src/exo/api/__init__.py b/src/exo/api/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/exo/master/adapters/__init__.py b/src/exo/api/adapters/__init__.py similarity index 100% rename from src/exo/master/adapters/__init__.py rename to src/exo/api/adapters/__init__.py diff --git a/src/exo/master/adapters/chat_completions.py b/src/exo/api/adapters/chat_completions.py similarity index 99% rename from src/exo/master/adapters/chat_completions.py rename to src/exo/api/adapters/chat_completions.py index c2221d71..2272cedc 100644 --- a/src/exo/master/adapters/chat_completions.py +++ b/src/exo/api/adapters/chat_completions.py @@ -4,7 +4,7 @@ import time from collections.abc import AsyncGenerator from typing import Any -from exo.shared.types.api import ( +from exo.api.types import ( ChatCompletionChoice, ChatCompletionMessage, ChatCompletionMessageText, diff --git a/src/exo/master/adapters/claude.py b/src/exo/api/adapters/claude.py similarity index 99% rename from src/exo/master/adapters/claude.py rename to src/exo/api/adapters/claude.py index 7a31f6fd..7ce94200 100644 --- a/src/exo/master/adapters/claude.py +++ b/src/exo/api/adapters/claude.py @@ -5,14 +5,8 @@ import re from collections.abc import AsyncGenerator from typing import Any -from exo.shared.types.api import FinishReason, Usage -from exo.shared.types.chunks import ( - ErrorChunk, - PrefillProgressChunk, - TokenChunk, - ToolCallChunk, -) -from exo.shared.types.claude_api import ( +from exo.api.types import FinishReason, Usage +from exo.api.types.claude_api import ( ClaudeContentBlock, ClaudeContentBlockDeltaEvent, ClaudeContentBlockStartEvent, @@ -35,6 +29,12 @@ from exo.shared.types.claude_api import ( ClaudeToolUseBlock, ClaudeUsage, ) +from exo.shared.types.chunks import ( + ErrorChunk, + PrefillProgressChunk, + TokenChunk, + ToolCallChunk, +) from exo.shared.types.common import CommandId from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams diff --git a/src/exo/master/adapters/ollama.py b/src/exo/api/adapters/ollama.py similarity index 99% rename from src/exo/master/adapters/ollama.py rename to src/exo/api/adapters/ollama.py index 75b772d0..3eeeb46c 100644 --- a/src/exo/master/adapters/ollama.py +++ b/src/exo/api/adapters/ollama.py @@ -4,14 +4,7 @@ import json from collections.abc import AsyncGenerator from typing import Any -from exo.shared.types.chunks import ( - ErrorChunk, - PrefillProgressChunk, - TokenChunk, - ToolCallChunk, -) -from exo.shared.types.common import CommandId -from exo.shared.types.ollama_api import ( +from exo.api.types.ollama_api import ( OllamaChatRequest, OllamaChatResponse, OllamaDoneReason, @@ -21,6 +14,13 @@ from exo.shared.types.ollama_api import ( OllamaToolCall, OllamaToolFunction, ) +from exo.shared.types.chunks import ( + ErrorChunk, + PrefillProgressChunk, + TokenChunk, + ToolCallChunk, +) +from exo.shared.types.common import CommandId from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams diff --git a/src/exo/master/adapters/responses.py b/src/exo/api/adapters/responses.py similarity index 99% rename from src/exo/master/adapters/responses.py rename to src/exo/api/adapters/responses.py index c92c4e37..014aeaf0 100644 --- a/src/exo/master/adapters/responses.py +++ b/src/exo/api/adapters/responses.py @@ -4,15 +4,8 @@ from collections.abc import AsyncGenerator from itertools import count from typing import Any -from exo.shared.types.api import Usage -from exo.shared.types.chunks import ( - ErrorChunk, - PrefillProgressChunk, - TokenChunk, - ToolCallChunk, -) -from exo.shared.types.common import CommandId -from exo.shared.types.openai_responses import ( +from exo.api.types import Usage +from exo.api.types.openai_responses import ( FunctionCallInputItem, ResponseCompletedEvent, ResponseContentPart, @@ -42,6 +35,13 @@ from exo.shared.types.openai_responses import ( ResponseTextDoneEvent, ResponseUsage, ) +from exo.shared.types.chunks import ( + ErrorChunk, + PrefillProgressChunk, + TokenChunk, + ToolCallChunk, +) +from exo.shared.types.common import CommandId from exo.shared.types.text_generation import ( InputMessage, TextGenerationTaskParams, diff --git a/src/exo/master/api.py b/src/exo/api/main.py similarity index 99% rename from src/exo/master/api.py rename to src/exo/api/main.py index 8c92fc21..08d95880 100644 --- a/src/exo/master/api.py +++ b/src/exo/api/main.py @@ -21,17 +21,17 @@ from hypercorn.config import Config from hypercorn.typing import ASGIFramework from loguru import logger -from exo.master.adapters.chat_completions import ( +from exo.api.adapters.chat_completions import ( chat_request_to_text_generation, collect_chat_response, generate_chat_stream, ) -from exo.master.adapters.claude import ( +from exo.api.adapters.claude import ( claude_request_to_text_generation, collect_claude_response, generate_claude_stream, ) -from exo.master.adapters.ollama import ( +from exo.api.adapters.ollama import ( collect_ollama_chat_response, collect_ollama_generate_response, generate_ollama_chat_stream, @@ -39,34 +39,12 @@ from exo.master.adapters.ollama import ( ollama_generate_request_to_text_generation, ollama_request_to_text_generation, ) -from exo.master.adapters.responses import ( +from exo.api.adapters.responses import ( collect_responses_response, generate_responses_stream, responses_request_to_text_generation, ) -from exo.master.event_log import DiskEventLog -from exo.master.image_store import ImageStore -from exo.master.placement import place_instance as get_instance_placements -from exo.shared.apply import apply -from exo.shared.constants import ( - DASHBOARD_DIR, - EXO_CACHE_HOME, - EXO_EVENT_LOG_DIR, - EXO_IMAGE_CACHE_DIR, - EXO_MAX_CHUNK_SIZE, - EXO_TRACING_CACHE_DIR, -) -from exo.shared.election import ElectionMessage -from exo.shared.logging import InterceptLogger -from exo.shared.models.model_cards import ( - ModelCard, - ModelId, - delete_custom_card, - get_model_cards, - is_custom_card, -) -from exo.shared.tracing import TraceEvent, compute_stats, export_trace, load_trace_file -from exo.shared.types.api import ( +from exo.api.types import ( AddCustomModelParams, AdvancedImageParams, BenchChatCompletionRequest, @@ -114,6 +92,48 @@ from exo.shared.types.api import ( TraceStatsResponse, normalize_image_size, ) +from exo.api.types.claude_api import ( + ClaudeMessagesRequest, + ClaudeMessagesResponse, +) +from exo.api.types.ollama_api import ( + OllamaChatRequest, + OllamaChatResponse, + OllamaGenerateRequest, + OllamaGenerateResponse, + OllamaModelDetails, + OllamaModelTag, + OllamaPsModel, + OllamaPsResponse, + OllamaShowRequest, + OllamaShowResponse, + OllamaTagsResponse, +) +from exo.api.types.openai_responses import ( + ResponsesRequest, + ResponsesResponse, +) +from exo.master.image_store import ImageStore +from exo.master.placement import place_instance as get_instance_placements +from exo.shared.apply import apply +from exo.shared.constants import ( + DASHBOARD_DIR, + EXO_CACHE_HOME, + EXO_EVENT_LOG_DIR, + EXO_IMAGE_CACHE_DIR, + EXO_MAX_CHUNK_SIZE, + EXO_TRACING_CACHE_DIR, +) +from exo.shared.election import ElectionMessage +from exo.shared.logging import InterceptLogger +from exo.shared.models.model_cards import ( + ModelCard, + ModelId, + delete_custom_card, + get_model_cards, + is_custom_card, +) +from exo.shared.tracing import TraceEvent, compute_stats, export_trace, load_trace_file from exo.shared.types.chunks import ( ErrorChunk, ImageChunk, @@ -122,10 +142,6 @@ from exo.shared.types.chunks import ( TokenChunk, ToolCallChunk, ) -from exo.shared.types.claude_api import ( - ClaudeMessagesRequest, - ClaudeMessagesResponse, -) from exo.shared.types.commands import ( Command, CreateInstance, @@ -151,29 +167,13 @@ from exo.shared.types.events import ( TracesMerged, ) from exo.shared.types.memory import Memory -from exo.shared.types.ollama_api import ( - OllamaChatRequest, - OllamaChatResponse, - OllamaGenerateRequest, - OllamaGenerateResponse, - OllamaModelDetails, - OllamaModelTag, - OllamaPsModel, - OllamaPsResponse, - OllamaShowRequest, - OllamaShowResponse, - OllamaTagsResponse, -) -from exo.shared.types.openai_responses import ( - ResponsesRequest, - ResponsesResponse, -) from exo.shared.types.state import State from exo.shared.types.worker.downloads import DownloadCompleted from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta from exo.shared.types.worker.shards import Sharding from exo.utils.banner import print_startup_banner from exo.utils.channels import Receiver, Sender, channel +from exo.utils.disk_event_log import DiskEventLog from exo.utils.power_sampler import PowerSampler from exo.utils.task_group import TaskGroup diff --git a/src/exo/master/tests/test_api_error_handling.py b/src/exo/api/tests/test_api_error_handling.py similarity index 97% rename from src/exo/master/tests/test_api_error_handling.py rename to src/exo/api/tests/test_api_error_handling.py index 631ba2a9..28cacf05 100644 --- a/src/exo/master/tests/test_api_error_handling.py +++ b/src/exo/api/tests/test_api_error_handling.py @@ -4,10 +4,11 @@ from typing import Any from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient +from exo.api.main import API + def test_http_exception_handler_formats_openai_style() -> None: """Test that HTTPException is converted to OpenAI-style error format.""" - from exo.master.api import API app = FastAPI() diff --git a/src/exo/master/tests/test_cancel_command.py b/src/exo/api/tests/test_cancel_command.py similarity index 98% rename from src/exo/master/tests/test_cancel_command.py rename to src/exo/api/tests/test_cancel_command.py index b64fc0cc..9da01408 100644 --- a/src/exo/master/tests/test_cancel_command.py +++ b/src/exo/api/tests/test_cancel_command.py @@ -5,12 +5,12 @@ from unittest.mock import AsyncMock, MagicMock from fastapi import FastAPI from fastapi.testclient import TestClient +from exo.api.main import API from exo.shared.types.common import CommandId def _make_api() -> Any: """Create a minimal API instance with cancel route and error handler.""" - from exo.master.api import API app = FastAPI() api = object.__new__(API) diff --git a/src/exo/master/tests/test_claude_api.py b/src/exo/api/tests/test_claude_api.py similarity index 98% rename from src/exo/master/tests/test_claude_api.py rename to src/exo/api/tests/test_claude_api.py index 08b2bef7..27e8cc4e 100644 --- a/src/exo/master/tests/test_claude_api.py +++ b/src/exo/api/tests/test_claude_api.py @@ -3,11 +3,11 @@ import pydantic import pytest -from exo.master.adapters.claude import ( +from exo.api.adapters.claude import ( claude_request_to_text_generation, finish_reason_to_claude_stop_reason, ) -from exo.shared.types.claude_api import ( +from exo.api.types.claude_api import ( ClaudeMessage, ClaudeMessagesRequest, ClaudeTextBlock, diff --git a/src/exo/master/tests/test_claude_tool_use.py b/src/exo/api/tests/test_claude_tool_use.py similarity index 99% rename from src/exo/master/tests/test_claude_tool_use.py rename to src/exo/api/tests/test_claude_tool_use.py index 3e814dce..3a2e505a 100644 --- a/src/exo/master/tests/test_claude_tool_use.py +++ b/src/exo/api/tests/test_claude_tool_use.py @@ -4,12 +4,12 @@ import json from collections.abc import AsyncGenerator from typing import Any, cast -from exo.master.adapters.claude import ( +from exo.api.adapters.claude import ( ClaudeMessagesResponse, collect_claude_response, generate_claude_stream, ) -from exo.shared.types.api import ToolCallItem +from exo.api.types import ToolCallItem from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk from exo.shared.types.common import CommandId, ModelId diff --git a/src/exo/master/tests/test_openai_responses_api.py b/src/exo/api/tests/test_openai_responses_api.py similarity index 96% rename from src/exo/master/tests/test_openai_responses_api.py rename to src/exo/api/tests/test_openai_responses_api.py index 78e3b741..7165df8e 100644 --- a/src/exo/master/tests/test_openai_responses_api.py +++ b/src/exo/api/tests/test_openai_responses_api.py @@ -7,11 +7,11 @@ The responses adapter converts it to TextGenerationTaskParams for the pipeline. import pydantic import pytest -from exo.shared.types.common import ModelId -from exo.shared.types.openai_responses import ( +from exo.api.types.openai_responses import ( ResponseInputMessage, ResponsesRequest, ) +from exo.shared.types.common import ModelId class TestResponsesRequestValidation: diff --git a/src/exo/api/types/__init__.py b/src/exo/api/types/__init__.py new file mode 100644 index 00000000..4c6218d0 --- /dev/null +++ b/src/exo/api/types/__init__.py @@ -0,0 +1,57 @@ +from .api import AddCustomModelParams as AddCustomModelParams +from .api import AdvancedImageParams as AdvancedImageParams +from .api import BenchChatCompletionRequest as BenchChatCompletionRequest +from .api import BenchChatCompletionResponse as BenchChatCompletionResponse +from .api import BenchImageGenerationResponse as BenchImageGenerationResponse +from .api import BenchImageGenerationTaskParams as BenchImageGenerationTaskParams +from .api import CancelCommandResponse as CancelCommandResponse +from .api import ChatCompletionChoice as ChatCompletionChoice +from .api import ChatCompletionMessage as ChatCompletionMessage +from .api import ChatCompletionMessageText as ChatCompletionMessageText +from .api import ChatCompletionRequest as ChatCompletionRequest +from .api import ChatCompletionResponse as ChatCompletionResponse +from .api import CompletionTokensDetails as CompletionTokensDetails +from .api import CreateInstanceParams as CreateInstanceParams +from .api import CreateInstanceResponse as CreateInstanceResponse +from .api import DeleteDownloadResponse as DeleteDownloadResponse +from .api import DeleteInstanceResponse as DeleteInstanceResponse +from .api import DeleteTracesRequest as DeleteTracesRequest +from .api import DeleteTracesResponse as DeleteTracesResponse +from .api import ErrorInfo as ErrorInfo +from .api import ErrorResponse as ErrorResponse +from .api import FinishReason as FinishReason +from .api import GenerationStats as GenerationStats +from .api import HuggingFaceSearchResult as HuggingFaceSearchResult +from .api import ImageData as ImageData +from .api import ImageEditsTaskParams as ImageEditsTaskParams +from .api import ImageGenerationResponse as ImageGenerationResponse +from .api import ImageGenerationStats as ImageGenerationStats +from .api import ImageGenerationTaskParams as ImageGenerationTaskParams +from .api import ImageListItem as ImageListItem +from .api import ImageListResponse as ImageListResponse +from .api import ImageSize as ImageSize +from .api import Logprobs as Logprobs +from .api import LogprobsContentItem as LogprobsContentItem +from .api import ModelList as ModelList +from .api import ModelListModel as ModelListModel +from .api import NodePowerStats as NodePowerStats +from .api import PlaceInstanceParams as PlaceInstanceParams +from .api import PlacementPreview as PlacementPreview +from .api import PlacementPreviewResponse as PlacementPreviewResponse +from .api import PowerUsage as PowerUsage +from .api import PromptTokensDetails as PromptTokensDetails +from .api import StartDownloadParams as StartDownloadParams +from .api import StartDownloadResponse as StartDownloadResponse +from .api import StreamingChoiceResponse as StreamingChoiceResponse +from .api import ToolCall as ToolCall +from .api import ToolCallItem as ToolCallItem +from .api import TopLogprobItem as TopLogprobItem +from .api import TraceCategoryStats as TraceCategoryStats +from .api import TraceEventResponse as TraceEventResponse +from .api import TraceListItem as TraceListItem +from .api import TraceListResponse as TraceListResponse +from .api import TraceRankStats as TraceRankStats +from .api import TraceResponse as TraceResponse +from .api import TraceStatsResponse as TraceStatsResponse +from .api import Usage as Usage +from .api import normalize_image_size as normalize_image_size diff --git a/src/exo/shared/types/api.py b/src/exo/api/types/api.py similarity index 100% rename from src/exo/shared/types/api.py rename to src/exo/api/types/api.py diff --git a/src/exo/shared/types/claude_api.py b/src/exo/api/types/claude_api.py similarity index 100% rename from src/exo/shared/types/claude_api.py rename to src/exo/api/types/claude_api.py diff --git a/src/exo/shared/types/ollama_api.py b/src/exo/api/types/ollama_api.py similarity index 100% rename from src/exo/shared/types/ollama_api.py rename to src/exo/api/types/ollama_api.py diff --git a/src/exo/shared/types/openai_responses.py b/src/exo/api/types/openai_responses.py similarity index 100% rename from src/exo/shared/types/openai_responses.py rename to src/exo/api/types/openai_responses.py diff --git a/src/exo/main.py b/src/exo/main.py index 2ecb62c2..3e08fa99 100644 --- a/src/exo/main.py +++ b/src/exo/main.py @@ -11,9 +11,9 @@ from loguru import logger from pydantic import PositiveInt import exo.routing.topics as topics +from exo.api.main import API from exo.download.coordinator import DownloadCoordinator from exo.download.impl_shard_downloader import exo_shard_downloader -from exo.master.api import API # TODO: should API be in master? from exo.master.main import Master from exo.routing.event_router import EventRouter from exo.routing.router import Router, get_node_id_keypair diff --git a/src/exo/master/main.py b/src/exo/master/main.py index b316f64d..dc347741 100644 --- a/src/exo/master/main.py +++ b/src/exo/master/main.py @@ -3,7 +3,6 @@ from datetime import datetime, timedelta, timezone import anyio from loguru import logger -from exo.master.event_log import DiskEventLog from exo.master.placement import ( add_instance_to_placements, cancel_unnecessary_downloads, @@ -61,6 +60,7 @@ from exo.shared.types.tasks import ( ) from exo.shared.types.worker.instances import InstanceId from exo.utils.channels import Receiver, Sender +from exo.utils.disk_event_log import DiskEventLog from exo.utils.event_buffer import MultiSourceBuffer from exo.utils.task_group import TaskGroup diff --git a/src/exo/shared/types/chunks.py b/src/exo/shared/types/chunks.py index 64481bab..9faf0cd6 100644 --- a/src/exo/shared/types/chunks.py +++ b/src/exo/shared/types/chunks.py @@ -1,18 +1,18 @@ from collections.abc import Generator from typing import Any, Literal -from exo.shared.models.model_cards import ModelId -from exo.shared.types.api import ( +from exo.api.types import ( + FinishReason, GenerationStats, ImageGenerationStats, + ToolCallItem, TopLogprobItem, Usage, ) +from exo.shared.models.model_cards import ModelId from exo.utils.pydantic_ext import TaggedModel -from .api import FinishReason from .common import CommandId -from .worker.runner_response import ToolCallItem class BaseChunk(TaggedModel): diff --git a/src/exo/shared/types/commands.py b/src/exo/shared/types/commands.py index e1c37288..fb55632e 100644 --- a/src/exo/shared/types/commands.py +++ b/src/exo/shared/types/commands.py @@ -1,10 +1,10 @@ from pydantic import Field -from exo.shared.models.model_cards import ModelCard, ModelId -from exo.shared.types.api import ( +from exo.api.types import ( ImageEditsTaskParams, ImageGenerationTaskParams, ) +from exo.shared.models.model_cards import ModelCard, ModelId from exo.shared.types.chunks import InputImageChunk from exo.shared.types.common import CommandId, NodeId, SystemId from exo.shared.types.text_generation import TextGenerationTaskParams diff --git a/src/exo/shared/types/tasks.py b/src/exo/shared/types/tasks.py index e2e6b5e5..e764ec33 100644 --- a/src/exo/shared/types/tasks.py +++ b/src/exo/shared/types/tasks.py @@ -2,7 +2,7 @@ from enum import Enum from pydantic import Field -from exo.shared.types.api import ( +from exo.api.types import ( ImageEditsTaskParams, ImageGenerationTaskParams, ) diff --git a/src/exo/shared/types/worker/runner_response.py b/src/exo/shared/types/worker/runner_response.py index 19118e15..27a7a3a0 100644 --- a/src/exo/shared/types/worker/runner_response.py +++ b/src/exo/shared/types/worker/runner_response.py @@ -1,7 +1,7 @@ from collections.abc import Generator from typing import Any, Literal -from exo.shared.types.api import ( +from exo.api.types import ( FinishReason, GenerationStats, ImageGenerationStats, diff --git a/src/exo/master/event_log.py b/src/exo/utils/disk_event_log.py similarity index 100% rename from src/exo/master/event_log.py rename to src/exo/utils/disk_event_log.py diff --git a/src/exo/utils/power_sampler.py b/src/exo/utils/power_sampler.py index b336ec9c..b8e985a1 100644 --- a/src/exo/utils/power_sampler.py +++ b/src/exo/utils/power_sampler.py @@ -5,7 +5,7 @@ from typing import final import anyio -from exo.shared.types.api import NodePowerStats, PowerUsage +from exo.api.types import NodePowerStats, PowerUsage from exo.shared.types.common import NodeId from exo.shared.types.profiling import SystemPerformanceProfile diff --git a/src/exo/master/tests/test_event_log.py b/src/exo/utils/tests/test_event_log.py similarity index 98% rename from src/exo/master/tests/test_event_log.py rename to src/exo/utils/tests/test_event_log.py index ffb4fb8e..0da1023a 100644 --- a/src/exo/master/tests/test_event_log.py +++ b/src/exo/utils/tests/test_event_log.py @@ -2,8 +2,8 @@ from pathlib import Path import pytest -from exo.master.event_log import DiskEventLog from exo.shared.types.events import TestEvent +from exo.utils.disk_event_log import DiskEventLog @pytest.fixture diff --git a/src/exo/utils/tests/test_power_sampler.py b/src/exo/utils/tests/test_power_sampler.py index e9c223ce..69f4ccee 100644 --- a/src/exo/utils/tests/test_power_sampler.py +++ b/src/exo/utils/tests/test_power_sampler.py @@ -3,7 +3,7 @@ from collections.abc import Mapping import anyio import pytest -from exo.shared.types.api import PowerUsage +from exo.api.types import PowerUsage from exo.shared.types.common import NodeId from exo.shared.types.profiling import SystemPerformanceProfile from exo.utils.power_sampler import PowerSampler diff --git a/src/exo/worker/engines/image/distributed_model.py b/src/exo/worker/engines/image/distributed_model.py index 88b60836..40a4fd38 100644 --- a/src/exo/worker/engines/image/distributed_model.py +++ b/src/exo/worker/engines/image/distributed_model.py @@ -6,8 +6,8 @@ import mlx.core as mx from mflux.models.common.config.config import Config from PIL import Image +from exo.api.types import AdvancedImageParams from exo.download.download_utils import build_model_path -from exo.shared.types.api import AdvancedImageParams from exo.shared.types.worker.instances import BoundInstance from exo.shared.types.worker.shards import CfgShardMetadata, PipelineShardMetadata from exo.worker.engines.image.config import ImageModelConfig diff --git a/src/exo/worker/engines/image/generate.py b/src/exo/worker/engines/image/generate.py index 2f4c5a4e..15523e51 100644 --- a/src/exo/worker/engines/image/generate.py +++ b/src/exo/worker/engines/image/generate.py @@ -9,7 +9,7 @@ from typing import Generator, Literal import mlx.core as mx from PIL import Image -from exo.shared.types.api import ( +from exo.api.types import ( AdvancedImageParams, ImageEditsTaskParams, ImageGenerationStats, diff --git a/src/exo/worker/engines/mlx/dsml_encoding.py b/src/exo/worker/engines/mlx/dsml_encoding.py index d5a5e210..8005ad49 100644 --- a/src/exo/worker/engines/mlx/dsml_encoding.py +++ b/src/exo/worker/engines/mlx/dsml_encoding.py @@ -4,7 +4,7 @@ from typing import Any from mlx_lm.chat_templates import deepseek_v32 -from exo.shared.types.api import ToolCallItem +from exo.api.types import ToolCallItem BOS_TOKEN: str = deepseek_v32.bos_token EOS_TOKEN: str = deepseek_v32.eos_token diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py index cb3ca5e7..3c8ef562 100644 --- a/src/exo/worker/engines/mlx/generator/batch_generate.py +++ b/src/exo/worker/engines/mlx/generator/batch_generate.py @@ -10,7 +10,7 @@ from mlx_lm.models.cache import RotatingKVCache from mlx_lm.sample_utils import make_logits_processors, make_sampler from mlx_lm.tokenizer_utils import StreamingDetokenizer, TokenizerWrapper -from exo.shared.types.api import ( +from exo.api.types import ( CompletionTokensDetails, FinishReason, GenerationStats, diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index eda6a7e5..f5470014 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -13,7 +13,7 @@ from mlx_lm.models.cache import ArraysCache, RotatingKVCache from mlx_lm.sample_utils import make_logits_processors, make_sampler from mlx_lm.tokenizer_utils import TokenizerWrapper -from exo.shared.types.api import ( +from exo.api.types import ( CompletionTokensDetails, FinishReason, GenerationStats, diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index f993a261..5389e78d 100644 --- a/src/exo/worker/main.py +++ b/src/exo/worker/main.py @@ -5,10 +5,10 @@ import anyio from anyio import fail_after from loguru import logger +from exo.api.types import ImageEditsTaskParams from exo.download.download_utils import resolve_model_in_path from exo.shared.apply import apply from exo.shared.models.model_cards import ModelId -from exo.shared.types.api import ImageEditsTaskParams from exo.shared.types.commands import ( ForwarderCommand, ForwarderDownloadCommand, diff --git a/src/exo/worker/runner/image_models/runner.py b/src/exo/worker/runner/image_models/runner.py index 385192cd..79eaf04f 100644 --- a/src/exo/worker/runner/image_models/runner.py +++ b/src/exo/worker/runner/image_models/runner.py @@ -4,10 +4,10 @@ from typing import TYPE_CHECKING, Literal import mlx.core as mx +from exo.api.types import ImageGenerationStats from exo.shared.constants import EXO_MAX_CHUNK_SIZE, EXO_TRACING_ENABLED from exo.shared.models.model_cards import ModelTask from exo.shared.tracing import clear_trace_buffer, get_trace_buffer -from exo.shared.types.api import ImageGenerationStats from exo.shared.types.chunks import ErrorChunk, ImageChunk from exo.shared.types.common import CommandId, ModelId from exo.shared.types.events import ( diff --git a/src/exo/worker/runner/llm_inference/model_output_parsers.py b/src/exo/worker/runner/llm_inference/model_output_parsers.py index cb5570e1..c6c7dd4e 100644 --- a/src/exo/worker/runner/llm_inference/model_output_parsers.py +++ b/src/exo/worker/runner/llm_inference/model_output_parsers.py @@ -13,7 +13,7 @@ from openai_harmony import ( # pyright: ignore[reportMissingTypeStubs] load_harmony_encoding, ) -from exo.shared.types.api import ToolCallItem +from exo.api.types import ToolCallItem from exo.shared.types.common import ModelId from exo.shared.types.mlx import Model from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse diff --git a/src/exo/worker/runner/llm_inference/tool_parsers.py b/src/exo/worker/runner/llm_inference/tool_parsers.py index 8fd02620..26140d9c 100644 --- a/src/exo/worker/runner/llm_inference/tool_parsers.py +++ b/src/exo/worker/runner/llm_inference/tool_parsers.py @@ -3,7 +3,7 @@ import math from dataclasses import dataclass from typing import Any, Callable -from exo.shared.types.api import ToolCallItem +from exo.api.types import ToolCallItem @dataclass diff --git a/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py b/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py index f83dcd36..ba331336 100644 --- a/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py +++ b/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py @@ -1,6 +1,6 @@ from collections.abc import Generator -from exo.shared.types.api import FinishReason +from exo.api.types import FinishReason from exo.shared.types.worker.runner_response import ( GenerationResponse, ToolCallResponse,