From 91c4d2b5fc4c8a0424a18cccc6bbe2ceb39f923a Mon Sep 17 00:00:00 2001 From: Andrei Cravtov Date: Thu, 11 Jun 2026 20:06:00 +0100 Subject: [PATCH] fix all tests and usages locator -> bootstrap --- src/exo/api/main.py | 10 +++---- src/exo/download/coordinator.py | 8 +++-- src/exo/download/download_utils.py | 10 +++---- .../tests/test_download_verification.py | 30 +++++++++---------- src/exo/download/tests/test_model_dirs.py | 20 ++++++++----- src/exo/download/tests/test_offline_mode.py | 10 +++---- src/exo/main.py | 2 +- src/exo/master/main.py | 4 +-- src/exo/routing/router.py | 5 ++-- src/exo/shared/models/model_cards.py | 8 +++-- ...cator_paths.py => test_bootstrap_paths.py} | 24 +++++++-------- src/exo/utils/banner.py | 4 +-- src/exo/utils/info_gatherer/info_gatherer.py | 4 +-- .../worker/engines/image/pipeline/runner.py | 3 +- src/exo/worker/runner/bootstrap.py | 2 +- src/exo/worker/runner/supervisor.py | 12 ++++---- .../tests/unittests/test_mlx/conftest.py | 4 +-- .../test_pipeline_prefill_callbacks.py | 5 ++-- .../test_prefix_cache_architectures.py | 4 +-- 19 files changed, 89 insertions(+), 80 deletions(-) rename src/exo/shared/tests/{test_locator_paths.py => test_bootstrap_paths.py} (90%) diff --git a/src/exo/api/main.py b/src/exo/api/main.py index 2ffe222db..ae90cfb97 100644 --- a/src/exo/api/main.py +++ b/src/exo/api/main.py @@ -23,6 +23,7 @@ from hypercorn.typing import ASGIFramework from hypercorn.utils import LifespanTimeoutError, ShutdownError from loguru import logger +import exo.shared.config as config from exo.api.adapters.chat_completions import ( chat_request_to_text_generation, collect_chat_response, @@ -126,7 +127,6 @@ from exo.api.types.openai_responses import ( 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.config import locator from exo.shared.constants import ( DASHBOARD_DIR, ENABLE_DISAGGREGATION, @@ -242,7 +242,7 @@ class API: election_receiver: Receiver[ElectionMessage], ) -> None: self.state = State() - self._api_event_log_dir = locator().event_log_dir / "api" + self._api_event_log_dir = config.bootstrap().event_log_dir / "api" self._event_log = DiskEventLog(self._api_event_log_dir) self._system_id = SystemId() self.command_sender = command_sender @@ -253,9 +253,9 @@ class API: self.last_completed_election: int = 0 self.port = port self._sent_image_hashes: set[str] = set() - self._tracing_cache_dir = locator().tracing_cache_dir + self._tracing_cache_dir = config.bootstrap().tracing_cache_dir self._onboarding_complete_file = ( - locator().exo_home.cache / "onboarding_complete" + config.bootstrap().exo_home.cache / "onboarding_complete" ) self.paused: bool = False @@ -291,7 +291,7 @@ class API: self._image_generation_queues: dict[ CommandId, Sender[ImageChunk | ErrorChunk] ] = {} - self._image_store = ImageStore(locator().image_cache_dir) + self._image_store = ImageStore(config.bootstrap().image_cache_dir) self._tg: TaskGroup = TaskGroup() def reset(self, result_clock: int, event_receiver: Receiver[IndexedEvent]): diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py index 393b5608e..571e1d9cb 100644 --- a/src/exo/download/coordinator.py +++ b/src/exo/download/coordinator.py @@ -7,6 +7,7 @@ import anyio from anyio import BrokenResourceError, ClosedResourceError, current_time, to_thread from loguru import logger +import exo.shared.config as config from exo.download.download_utils import ( RepoDownloadProgress, delete_model, @@ -19,7 +20,6 @@ from exo.routing.event_router import ( EventRouterBrokenResourceError, EventRouterClosedResourceError, ) -from exo.shared.config import locator from exo.shared.models import model_cards from exo.shared.models.model_cards import ModelId from exo.shared.types.commands import ( @@ -69,7 +69,9 @@ class DownloadCoordinator: @staticmethod def _default_model_dir(model_id: ModelId) -> str: - return str(locator().models_dirs.default_models_dir / model_id.normalize()) + return str( + config.bootstrap().models_dirs.default_models_dir / model_id.normalize() + ) def _completed_from_path( self, @@ -433,7 +435,7 @@ class DownloadCoordinator: NodeDownloadProgress(download_progress=status) ) # Scan read-only directories for pre-downloaded models - if locator().models_dirs.models_read_only_dirs: + if config.bootstrap().models_dirs.models_read_only_dirs: for card in await model_cards.card_cache.list_all(): mid = card.model_id if mid in self.active_downloads: diff --git a/src/exo/download/download_utils.py b/src/exo/download/download_utils.py index 678cb30d8..35abdf9e9 100644 --- a/src/exo/download/download_utils.py +++ b/src/exo/download/download_utils.py @@ -24,6 +24,7 @@ from pydantic import ( TypeAdapter, ) +import exo.shared.config as config from exo.download.huggingface_utils import ( filter_repo_objects, get_allow_patterns, @@ -31,7 +32,6 @@ from exo.download.huggingface_utils import ( get_hf_endpoint, get_hf_token, ) -from exo.shared.config import locator from exo.shared.models.model_cards import ModelCard, ModelTask from exo.shared.types.common import ModelId from exo.shared.types.memory import Memory @@ -46,19 +46,19 @@ from exo.shared.types.worker.shards import ShardMetadata def _default_models_dir() -> Path: - return locator().models_dirs.default_models_dir + return config.bootstrap().models_dirs.default_models_dir def _writable_models_dirs() -> list[Path]: - return locator().models_dirs.models_dirs + return config.bootstrap().models_dirs.models_dirs def _read_only_models_dirs() -> list[Path]: - return locator().models_dirs.models_read_only_dirs + return config.bootstrap().models_dirs.models_read_only_dirs def _model_search_dirs() -> tuple[Path, ...]: - models_dirs = locator().models_dirs + models_dirs = config.bootstrap().models_dirs return (*models_dirs.models_read_only_dirs, *models_dirs.models_dirs) diff --git a/src/exo/download/tests/test_download_verification.py b/src/exo/download/tests/test_download_verification.py index 35dd12ad5..698758ef8 100644 --- a/src/exo/download/tests/test_download_verification.py +++ b/src/exo/download/tests/test_download_verification.py @@ -8,7 +8,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import aiofiles import aiofiles.os as aios import pytest -from exo_rs import LocatorConfig +from exo_rs import BootstrapSettings from pydantic import TypeAdapter from exo.download.download_utils import ( @@ -25,8 +25,8 @@ def model_id() -> ModelId: return ModelId("test-org/test-model") -def _mock_locator_config(models_dir: Path) -> LocatorConfig: - cfg = LocatorConfig.default() +def _mock_bootstrap_settings(models_dir: Path) -> BootstrapSettings: + cfg = BootstrapSettings.default() models_dirs = cfg.models_dirs models_dirs.default_models_dir = models_dir models_dirs.models_dirs = [models_dir] @@ -187,11 +187,11 @@ class TestFileListCache: FileListEntry(type="file", path="config.json", size=100), ] - cfg = _mock_locator_config(models_dir) + cfg = _mock_bootstrap_settings(models_dir) with ( patch( - "exo.download.download_utils.locator", + "exo.download.download_utils.config.bootstrap", new_callable=Mock, return_value=cfg, ), @@ -239,11 +239,11 @@ class TestFileListCache: TypeAdapter(list[FileListEntry]).dump_json(cached_file_list).decode() ) - cfg = _mock_locator_config(models_dir) + cfg = _mock_bootstrap_settings(models_dir) with ( patch( - "exo.download.download_utils.locator", + "exo.download.download_utils.config.bootstrap", new_callable=Mock, return_value=cfg, ), @@ -263,11 +263,11 @@ class TestFileListCache: """Test that errors propagate when fetch fails and no cache exists.""" models_dir = tmp_path / "models" - cfg = _mock_locator_config(models_dir) + cfg = _mock_bootstrap_settings(models_dir) with ( patch( - "exo.download.download_utils.locator", + "exo.download.download_utils.config.bootstrap", new_callable=Mock, return_value=cfg, ), @@ -302,11 +302,11 @@ class TestModelDeletion: async with aiofiles.open(cache_dir / "file_list.json", "w") as f: await f.write("[]") - cfg = _mock_locator_config(models_dir) + cfg = _mock_bootstrap_settings(models_dir) with ( patch( - "exo.download.download_utils.locator", + "exo.download.download_utils.config.bootstrap", new_callable=Mock, return_value=cfg, ), @@ -329,11 +329,11 @@ class TestModelDeletion: async with aiofiles.open(cache_dir / "file_list.json", "w") as f: await f.write("[]") - cfg = _mock_locator_config(models_dir) + cfg = _mock_bootstrap_settings(models_dir) with ( patch( - "exo.download.download_utils.locator", + "exo.download.download_utils.config.bootstrap", new_callable=Mock, return_value=cfg, ), @@ -352,11 +352,11 @@ class TestModelDeletion: models_dir = tmp_path / "models" await aios.makedirs(models_dir, exist_ok=True) - cfg = _mock_locator_config(models_dir) + cfg = _mock_bootstrap_settings(models_dir) with ( patch( - "exo.download.download_utils.locator", + "exo.download.download_utils.config.bootstrap", new_callable=Mock, return_value=cfg, ), diff --git a/src/exo/download/tests/test_model_dirs.py b/src/exo/download/tests/test_model_dirs.py index 44dfe58d7..7daf1d434 100644 --- a/src/exo/download/tests/test_model_dirs.py +++ b/src/exo/download/tests/test_model_dirs.py @@ -10,7 +10,7 @@ from unittest.mock import Mock, patch import aiofiles import aiofiles.os as aios import pytest -from exo_rs import LocatorConfig +from exo_rs import BootstrapSettings from exo.download.download_utils import ( InsufficientDiskSpaceError, @@ -44,13 +44,13 @@ def _create_incomplete_model(model_dir: Path) -> None: # model.safetensors is missing -def _mock_locator_config( +def _mock_bootstrap_settings( default: Path, *, writable: list[Path] | None = None, read_only: list[Path] | None = None, -) -> LocatorConfig: - cfg = LocatorConfig.default() +) -> BootstrapSettings: + cfg = BootstrapSettings.default() models_dirs = cfg.models_dirs models_dirs.default_models_dir = default models_dirs.models_dirs = writable or [] @@ -66,11 +66,13 @@ def _patched_model_dirs( writable: Sequence[Path] = (), read_only: Sequence[Path] = (), ) -> Iterator[None]: - cfg = _mock_locator_config( + cfg = _mock_bootstrap_settings( default, writable=list(writable), read_only=list(read_only) ) with patch( - "exo.download.download_utils.locator", new_callable=Mock, return_value=cfg + "exo.download.download_utils.config.bootstrap", + new_callable=Mock, + return_value=cfg, ): yield @@ -256,10 +258,12 @@ class TestDeleteModel: await aios.makedirs(writable2, exist_ok=True) await aios.makedirs(default, exist_ok=True) - cfg = _mock_locator_config(default, writable=[writable1, writable2, default]) + cfg = _mock_bootstrap_settings( + default, writable=[writable1, writable2, default] + ) with patch( - "exo.download.download_utils.locator", + "exo.download.download_utils.config.bootstrap", new_callable=Mock, return_value=cfg, ): diff --git a/src/exo/download/tests/test_offline_mode.py b/src/exo/download/tests/test_offline_mode.py index de9cb24f3..2bfbb5aca 100644 --- a/src/exo/download/tests/test_offline_mode.py +++ b/src/exo/download/tests/test_offline_mode.py @@ -9,7 +9,7 @@ from unittest.mock import AsyncMock, Mock, patch import aiofiles import aiofiles.os as aios import pytest -from exo_rs import LocatorConfig +from exo_rs import BootstrapSettings from exo.download.download_utils import ( _download_file, # pyright: ignore[reportPrivateUsage] @@ -25,8 +25,8 @@ def model_id() -> ModelId: return ModelId("test-org/test-model") -def _mock_locator_config(models_dir: Path) -> LocatorConfig: - cfg = LocatorConfig.default() +def _mock_bootstrap_settings(models_dir: Path) -> BootstrapSettings: + cfg = BootstrapSettings.default() models_dirs = cfg.models_dirs models_dirs.default_models_dir = models_dir models_dirs.models_dirs = [models_dir] @@ -39,11 +39,11 @@ async def temp_models_dir(tmp_path: Path) -> AsyncIterator[Path]: models_dir = tmp_path / "models" await aios.makedirs(models_dir, exist_ok=True) - cfg = _mock_locator_config(models_dir) + cfg = _mock_bootstrap_settings(models_dir) with ( patch( - "exo.download.download_utils.locator", + "exo.download.download_utils.config.bootstrap", new_callable=Mock, return_value=cfg, ), diff --git a/src/exo/main.py b/src/exo/main.py index 97b1ca69b..cbf211134 100644 --- a/src/exo/main.py +++ b/src/exo/main.py @@ -15,6 +15,7 @@ from loguru import logger from pydantic import PositiveInt import exo.routing.topics as topics +import exo.shared.config as config from exo import __version__ from exo.api.main import API from exo.download.coordinator import DownloadCoordinator @@ -22,7 +23,6 @@ from exo.download.impl_shard_downloader import exo_shard_downloader from exo.master.main import Master from exo.routing.event_router import EventRouter from exo.routing.router import Router, get_node_zid -from exo.shared import config from exo.shared.election import Election, ElectionResult from exo.shared.logging import logger_cleanup, logger_setup from exo.shared.types.common import NodeId, SessionId diff --git a/src/exo/master/main.py b/src/exo/master/main.py index 2c0883ee6..d727ffe39 100644 --- a/src/exo/master/main.py +++ b/src/exo/master/main.py @@ -3,6 +3,7 @@ from datetime import datetime, timedelta, timezone import anyio from loguru import logger +import exo.shared.config as config from exo.master.placement import ( add_instance_to_placements, cancel_unnecessary_downloads, @@ -16,7 +17,6 @@ from exo.routing.event_router import ( EventRouterClosedResourceError, ) from exo.shared.apply import apply -from exo.shared.config import locator from exo.shared.constants import EXO_TRACING_ENABLED from exo.shared.types.commands import ( AddCustomModelCard, @@ -144,7 +144,7 @@ class Master: self.event_sender = event_sender self._system_id = SystemId() self._multi_buffer = MultiSourceBuffer[SystemId, Event]() - self._master_event_log_dir = locator().event_log_dir / "master" + self._master_event_log_dir = config.bootstrap().event_log_dir / "master" self._event_log = DiskEventLog(self._master_event_log_dir) self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {} self._expected_ranks: dict[TaskId, set[int]] = {} diff --git a/src/exo/routing/router.py b/src/exo/routing/router.py index cf25b865c..7ae19dcc5 100644 --- a/src/exo/routing/router.py +++ b/src/exo/routing/router.py @@ -17,7 +17,7 @@ from exo_rs import ( ) from loguru import logger -from exo.shared.config import locator +import exo.shared.config as config from exo.shared.types.common import NodeId from exo.utils.channels import Receiver, Sender, channel from exo.utils.pydantic_ext import FrozenModel @@ -236,8 +236,7 @@ def get_node_zid( Obtains the :class:`Keypair` associated with this node-ID. Obtain the :class:`PeerId` by from it. """ - if path is None: - path = locator().node_zid + path = path or config.bootstrap().node_zid # TODO(evan): bring back node id persistence once we figure out how to deal with duplicates return NodeId(os.urandom(16).hex().lstrip("0")) diff --git a/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py index 2bb87825b..7e57a0825 100644 --- a/src/exo/shared/models/model_cards.py +++ b/src/exo/shared/models/model_cards.py @@ -20,7 +20,7 @@ from pydantic import ( ) from tomlkit.exceptions import TOMLKitError -from exo.shared.config import locator +import exo.shared.config as config from exo.shared.constants import ( EXO_ENABLE_IMAGE_MODELS, RESOURCES_DIR, @@ -40,7 +40,7 @@ _BUILTIN_CARD_DIRS = [ def _custom_cards_dir() -> Path: - return Path(str(locator().custom_model_cards_dir)) + return Path(str(config.bootstrap().custom_model_cards_dir)) class _CardCache: @@ -99,7 +99,9 @@ card_cache = _CardCache() def detect_vision_from_config(model_id: ModelId) -> "VisionCardConfig | None": normalized = model_id.normalize() - for model_dir in [d / normalized for d in locator().models_dirs.models_dirs]: + for model_dir in [ + d / normalized for d in config.bootstrap().models_dirs.models_dirs + ]: config_path = model_dir / "config.json" if not config_path.exists(): continue diff --git a/src/exo/shared/tests/test_locator_paths.py b/src/exo/shared/tests/test_bootstrap_paths.py similarity index 90% rename from src/exo/shared/tests/test_locator_paths.py rename to src/exo/shared/tests/test_bootstrap_paths.py index aece05c09..f72d56941 100644 --- a/src/exo/shared/tests/test_locator_paths.py +++ b/src/exo/shared/tests/test_bootstrap_paths.py @@ -1,10 +1,10 @@ -"""Tests for XDG Base Directory Specification compliance.""" +"""Tests for bootstrap path and model-directory resolution.""" import sys from pathlib import Path import pytest -from exo_rs import LocatorConfig +from exo_rs import BootstrapSettings @pytest.mark.skipif(sys.platform != "linux", reason="XDG dirs are Linux-specific") @@ -19,7 +19,7 @@ def test_xdg_paths_on_linux(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("XDG_DATA_HOME", str(data_home)) monkeypatch.setenv("XDG_CACHE_HOME", str(cache_home)) - exo_home = LocatorConfig.from_env_only().exo_home + exo_home = BootstrapSettings.from_env_only().exo_home assert config_home / "exo" == exo_home.config assert data_home / "exo" == exo_home.data @@ -37,7 +37,7 @@ def test_standard_directories_on_macos(tmp_path: Path, monkeypatch: pytest.Monke monkeypatch.setenv("XDG_DATA_HOME", str(tmp_path / "ignored-data")) monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "ignored-cache")) - exo_home = LocatorConfig.from_env_only().exo_home + exo_home = BootstrapSettings.from_env_only().exo_home assert home / "Library" / "Application Support" / "exo" == exo_home.config assert home / "Library" / "Application Support" / "exo" == exo_home.data @@ -55,7 +55,7 @@ def test_xdg_default_paths_on_linux(tmp_path: Path, monkeypatch: pytest.MonkeyPa monkeypatch.delenv("XDG_CACHE_HOME", raising=False) monkeypatch.setenv("HOME", str(home)) - exo_home = LocatorConfig.from_env_only().exo_home + exo_home = BootstrapSettings.from_env_only().exo_home assert home / ".config" / "exo" == exo_home.config assert home / ".local" / "share" / "exo" == exo_home.data @@ -73,7 +73,7 @@ def test_legacy_exo_home_takes_precedence( monkeypatch.setenv("XDG_DATA_HOME", str(tmp_path / "ignored-data")) monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "ignored-cache")) - exo_home = LocatorConfig.from_env_only().exo_home + exo_home = BootstrapSettings.from_env_only().exo_home assert exo_home_path == exo_home.config assert exo_home_path == exo_home.data @@ -90,7 +90,7 @@ def test_models_in_data_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): monkeypatch.delenv("EXO_MODELS_DIRS", raising=False) monkeypatch.setenv("XDG_DATA_HOME", str(data_home)) - cfg = LocatorConfig.from_env_only() + cfg = BootstrapSettings.from_env_only() assert cfg.models_dirs.default_models_dir.parent == cfg.exo_home.data @@ -106,7 +106,7 @@ def test_default_dir_always_prepended_to_models_dirs( monkeypatch.delenv("EXO_MODELS_READ_ONLY_DIRS", raising=False) monkeypatch.setenv("EXO_MODELS_DIRS", str(custom_models_dir)) - models_dirs = LocatorConfig.from_env_only().models_dirs + models_dirs = BootstrapSettings.from_env_only().models_dirs assert models_dirs.models_dirs[0] == models_dirs.default_models_dir assert custom_models_dir in models_dirs.models_dirs @@ -121,7 +121,7 @@ def test_default_models_dir_override(tmp_path: Path, monkeypatch: pytest.MonkeyP monkeypatch.delenv("EXO_MODELS_READ_ONLY_DIRS", raising=False) monkeypatch.setenv("EXO_DEFAULT_MODELS_DIR", str(default_models_dir)) - models_dirs = LocatorConfig.from_env_only().models_dirs + models_dirs = BootstrapSettings.from_env_only().models_dirs assert default_models_dir == models_dirs.default_models_dir assert models_dirs.models_dirs[0] == models_dirs.default_models_dir @@ -134,7 +134,7 @@ def test_default_dir_only_entry_when_env_unset(monkeypatch: pytest.MonkeyPatch): monkeypatch.delenv("EXO_MODELS_DIRS", raising=False) monkeypatch.delenv("EXO_MODELS_READ_ONLY_DIRS", raising=False) - models_dirs = LocatorConfig.from_env_only().models_dirs + models_dirs = BootstrapSettings.from_env_only().models_dirs assert models_dirs.models_dirs == [models_dirs.default_models_dir] @@ -154,7 +154,7 @@ def test_overlap_between_dirs_and_read_only_dirs( monkeypatch.setenv("EXO_MODELS_DIRS", f"{shared}:{writable_only}") monkeypatch.setenv("EXO_MODELS_READ_ONLY_DIRS", f"{shared}:{read_only}") - models_dirs = LocatorConfig.from_env_only().models_dirs + models_dirs = BootstrapSettings.from_env_only().models_dirs assert shared not in models_dirs.models_dirs assert writable_only in models_dirs.models_dirs @@ -169,6 +169,6 @@ def test_empty_read_only_dirs_when_unset(monkeypatch: pytest.MonkeyPatch): monkeypatch.delenv("EXO_MODELS_DIRS", raising=False) monkeypatch.delenv("EXO_MODELS_READ_ONLY_DIRS", raising=False) - models_dirs = LocatorConfig.from_env_only().models_dirs + models_dirs = BootstrapSettings.from_env_only().models_dirs assert models_dirs.models_read_only_dirs == [] diff --git a/src/exo/utils/banner.py b/src/exo/utils/banner.py index 39b5c6a20..af9a683d8 100644 --- a/src/exo/utils/banner.py +++ b/src/exo/utils/banner.py @@ -4,13 +4,13 @@ import sys import webbrowser from pathlib import Path -from exo.shared.config import locator +import exo.shared.config as config logger = logging.getLogger(__name__) def _first_run_marker() -> Path: - return locator().exo_home.config / ".dashboard_opened" + return config.bootstrap().exo_home.config / ".dashboard_opened" def _is_first_run() -> bool: diff --git a/src/exo/utils/info_gatherer/info_gatherer.py b/src/exo/utils/info_gatherer/info_gatherer.py index d16e46b67..ad019e548 100644 --- a/src/exo/utils/info_gatherer/info_gatherer.py +++ b/src/exo/utils/info_gatherer/info_gatherer.py @@ -12,7 +12,7 @@ from anyio.streams.buffered import BufferedByteReceiveStream from loguru import logger from pydantic import ValidationError -from exo.shared.config import locator +import exo.shared.config as config from exo.shared.types.backends import Backend from exo.shared.types.memory import Memory from exo.shared.types.profiling import ( @@ -310,7 +310,7 @@ class NodeDiskUsage(TaggedModel): async def gather(cls) -> Self: return cls( disk_usage=await to_thread.run_sync( - DiskUsage.from_path, locator().models_dirs.default_models_dir + DiskUsage.from_path, config.bootstrap().models_dirs.default_models_dir ) ) diff --git a/src/exo/worker/engines/image/pipeline/runner.py b/src/exo/worker/engines/image/pipeline/runner.py index 2390d9912..d29375ad3 100644 --- a/src/exo/worker/engines/image/pipeline/runner.py +++ b/src/exo/worker/engines/image/pipeline/runner.py @@ -544,8 +544,7 @@ class DiffusionRunner: num_sync_steps: int, capture_steps: set[int] | None = None, ): - if capture_steps is None: - capture_steps = set() + capture_steps = capture_steps or set() self._reset_all_caches() clear_trace_buffer() diff --git a/src/exo/worker/runner/bootstrap.py b/src/exo/worker/runner/bootstrap.py index 9f060e06f..52e5c3cd6 100644 --- a/src/exo/worker/runner/bootstrap.py +++ b/src/exo/worker/runner/bootstrap.py @@ -7,7 +7,7 @@ from typing import Self, cast import loguru from exo_rs import AppSettings, BootstrapSettings -from exo.shared import config +import exo.shared.config as config from exo.shared.types.events import Event from exo.shared.types.tasks import Task, TaskId from exo.shared.types.worker.instances import BoundInstance diff --git a/src/exo/worker/runner/supervisor.py b/src/exo/worker/runner/supervisor.py index 55ba0efa6..c065f04b5 100644 --- a/src/exo/worker/runner/supervisor.py +++ b/src/exo/worker/runner/supervisor.py @@ -14,7 +14,7 @@ from anyio import ( ) from loguru import logger -from exo.shared import config +import exo.shared.config as config from exo.shared.types.chunks import ErrorChunk from exo.shared.types.events import ( ChunkGenerated, @@ -79,10 +79,12 @@ class RunnerStdioHandler: stdout_log_path: PathLike[str] | None = None, stderr_log_path: PathLike[str] | None = None, ) -> Self: - if stdout_log_path is None: - stdout_log_path = config.bootstrap().log_files.exo_runner_stdout_log - if stderr_log_path is None: - stderr_log_path = config.bootstrap().log_files.exo_runner_stderr_log + stdout_log_path = ( + stdout_log_path or config.bootstrap().log_files.exo_runner_stdout_log + ) + stderr_log_path = ( + stderr_log_path or config.bootstrap().log_files.exo_runner_stderr_log + ) # these are append only logs used to gather data for log template mining # diff --git a/src/exo/worker/tests/unittests/test_mlx/conftest.py b/src/exo/worker/tests/unittests/test_mlx/conftest.py index 1e9b74243..f2e9c18ee 100644 --- a/src/exo/worker/tests/unittests/test_mlx/conftest.py +++ b/src/exo/worker/tests/unittests/test_mlx/conftest.py @@ -10,7 +10,7 @@ from typing import Any, cast import mlx.core as mx import mlx.nn as nn -from exo.shared.config import locator +import exo.shared.config as config from exo.shared.models.model_cards import ModelCard, ModelTask from exo.shared.types.backends import Backend from exo.shared.types.common import ModelId @@ -53,7 +53,7 @@ def create_hostfile(world_size: int, base_port: int) -> tuple[str, list[str]]: # Use GPT OSS 20b to test as it is a model with a lot of strange behaviour DEFAULT_GPT_OSS_CONFIG = PipelineTestConfig( - model_path=locator().models_dirs.default_models_dir + model_path=config.bootstrap().models_dirs.default_models_dir / "mlx-community--gpt-oss-20b-MXFP4-Q8", total_layers=24, base_port=29600, diff --git a/src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.py b/src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.py index a65bfa051..2213853f1 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.py @@ -15,7 +15,7 @@ from typing import Any, cast import pytest -from exo.shared.config import locator +import exo.shared.config as config from exo.shared.models.model_cards import ModelCard, ModelTask from exo.shared.types.backends import Backend from exo.shared.types.common import ModelId @@ -24,7 +24,8 @@ from exo.shared.types.text_generation import InputMessage, TextGenerationTaskPar MODEL_ID = "mlx-community/gpt-oss-20b-MXFP4-Q8" MODEL_PATH = ( - locator().models_dirs.default_models_dir / "mlx-community--gpt-oss-20b-MXFP4-Q8" + config.bootstrap().models_dirs.default_models_dir + / "mlx-community--gpt-oss-20b-MXFP4-Q8" ) TOTAL_LAYERS = 24 MAX_TOKENS = 10 diff --git a/src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py b/src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py index 1331ac636..c5d55ac13 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py @@ -14,8 +14,8 @@ import pytest from mlx.utils import tree_flatten, tree_unflatten from mlx_lm.tokenizer_utils import TokenizerWrapper +import exo.shared.config as config from exo.download.download_utils import resolve_existing_model -from exo.shared.config import locator from exo.shared.types.common import ModelId from exo.shared.types.text_generation import ( InputMessage, @@ -107,7 +107,7 @@ def _reduce_config(cfg: dict[str, Any]) -> dict[str, Any]: def _search_dirs(): - m = locator().models_dirs + m = config.bootstrap().models_dirs return (*m.models_read_only_dirs, *m.models_dirs)