fix all tests and usages locator -> bootstrap

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]):

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:

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)
 
 

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,
             ),

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,
         ):

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,
         ),

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

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]] = {}

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"))

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

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 == []

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:

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
             )
         )
 

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()

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

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
         #

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,

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

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)