diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py index 0be66a1b..a5508f92 100644 --- a/src/exo/download/coordinator.py +++ b/src/exo/download/coordinator.py @@ -207,13 +207,9 @@ class DownloadCoordinator: except (FileNotFoundError, TOMLKitError, UnicodeDecodeError): pass - if config.max_storage is not None: - doc["max_storage_gb"] = round(config.max_storage.in_gb, 2) - else: - doc.pop("max_storage_gb", None) # pyright: ignore[reportUnknownMemberType] - doc.pop("max_storage_bytes", None) # pyright: ignore[reportUnknownMemberType] - - doc["storage_policy"] = config.storage_policy + # Clear max_storage_gb so it doesn't linger when max_storage is None + doc.pop("max_storage_gb", None) # pyright: ignore[reportUnknownMemberType] + doc.update(config.to_disk()) # pyright: ignore[reportUnknownMemberType] await cfg_path.write_text(tomlkit.dumps(doc)) # pyright: ignore[reportUnknownMemberType] logger.debug(f"Persisted storage config to {cfg_path}") diff --git a/src/exo/main.py b/src/exo/main.py index 5873ce9b..33e0ea0e 100644 --- a/src/exo/main.py +++ b/src/exo/main.py @@ -175,24 +175,17 @@ class Node: @staticmethod async def _load_storage_config(args: "Args") -> StorageConfig: """Load storage config: start from config.toml, overlay CLI args.""" - # Start from config.toml as the base - base_max_storage: Memory | None = None - base_policy: StoragePolicy = "manual" - node_config = await NodeConfig.gather() - if node_config is not None: - if node_config.max_storage_bytes is not None: - base_max_storage = Memory.from_bytes(node_config.max_storage_bytes) - base_policy = node_config.storage_policy + base = node_config.storage_config if node_config is not None else StorageConfig() # CLI args override individual fields (non-default values only) max_storage = ( Memory.from_gb(args.max_storage_gb) if args.max_storage_gb is not None - else base_max_storage + else base.max_storage ) storage_policy = ( - args.storage_policy if args.storage_policy != "manual" else base_policy + args.storage_policy if args.storage_policy != "manual" else base.storage_policy ) return StorageConfig(max_storage=max_storage, storage_policy=storage_policy) diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py index 3d059890..ab6ca4fb 100644 --- a/src/exo/shared/apply.py +++ b/src/exo/shared/apply.py @@ -30,7 +30,6 @@ from exo.shared.types.events import ( TracesCollected, TracesMerged, ) -from exo.shared.types.memory import Memory from exo.shared.types.profiling import ( NodeIdentity, NodeNetworkInfo, @@ -39,7 +38,6 @@ from exo.shared.types.profiling import ( ThunderboltBridgeStatus, ) from exo.shared.types.state import State -from exo.shared.types.storage import StorageConfig from exo.shared.types.tasks import Task, TaskId, TaskStatus from exo.shared.types.topology import Connection, RDMAConnection from exo.shared.types.worker.downloads import DownloadProgress @@ -304,17 +302,9 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State: case NodeDiskUsage(): update["node_disk"] = {**state.node_disk, event.node_id: info.disk_usage} case NodeConfig(): - storage_config = StorageConfig( - max_storage=( - Memory.from_bytes(info.max_storage_bytes) - if info.max_storage_bytes is not None - else None - ), - storage_policy=info.storage_policy, - ) update["node_storage_config"] = { **state.node_storage_config, - event.node_id: storage_config, + event.node_id: info.storage_config, } case MiscData(): current_identity = state.node_identities.get(event.node_id, NodeIdentity()) diff --git a/src/exo/shared/types/storage.py b/src/exo/shared/types/storage.py index e42b29c8..54fb16f8 100644 --- a/src/exo/shared/types/storage.py +++ b/src/exo/shared/types/storage.py @@ -1,4 +1,4 @@ -from typing import Literal, final +from typing import Any, Literal, Self, final from exo.shared.types.memory import Memory from exo.utils.pydantic_ext import FrozenModel @@ -10,3 +10,20 @@ StoragePolicy = Literal["manual", "auto-evict"] class StorageConfig(FrozenModel): max_storage: Memory | None = None storage_policy: StoragePolicy = "manual" + + @classmethod + def from_disk(cls, data: dict[str, Any]) -> Self: + """Parse from a TOML config dict (e.g. from tomllib).""" + max_storage: Memory | None = None + if "max_storage_gb" in data: + max_storage = Memory.from_gb(float(data["max_storage_gb"])) # pyright: ignore[reportAny] + policy: StoragePolicy = data.get("storage_policy", "manual") # pyright: ignore[reportAny] + return cls(max_storage=max_storage, storage_policy=policy) + + def to_disk(self) -> dict[str, Any]: + """Serialize to a dict suitable for writing to TOML.""" + result: dict[str, Any] = {} + if self.max_storage is not None: + result["max_storage_gb"] = round(self.max_storage.in_gb, 2) + result["storage_policy"] = self.storage_policy + return result diff --git a/src/exo/utils/info_gatherer/info_gatherer.py b/src/exo/utils/info_gatherer/info_gatherer.py index 764d401e..4a594e02 100644 --- a/src/exo/utils/info_gatherer/info_gatherer.py +++ b/src/exo/utils/info_gatherer/info_gatherer.py @@ -15,13 +15,13 @@ from pydantic import ValidationError from exo.shared.constants import EXO_CONFIG_FILE, EXO_DEFAULT_MODELS_DIR from exo.shared.types.memory import Memory +from exo.shared.types.storage import StorageConfig from exo.shared.types.profiling import ( DiskUsage, MemoryUsage, NetworkInterfaceInfo, ThunderboltBridgeStatus, ) -from exo.shared.types.storage import StoragePolicy from exo.shared.types.thunderbolt import ( ThunderboltConnection, ThunderboltConnectivity, @@ -295,8 +295,7 @@ class ThunderboltBridgeInfo(TaggedModel): class NodeConfig(TaggedModel): """Node configuration from EXO_CONFIG_FILE, reloaded from the file only at startup. Other changes should come in through the API and propagate from there""" - max_storage_bytes: int | None = None - storage_policy: StoragePolicy = "manual" + storage_config: StorageConfig = StorageConfig() @classmethod async def gather(cls) -> Self | None: @@ -307,10 +306,7 @@ class NodeConfig(TaggedModel): try: contents = (await f.read()).decode("utf-8") data = tomllib.loads(contents) - if "max_storage_gb" in data: - gb_value: float = float(data.pop("max_storage_gb")) # pyright: ignore[reportAny] - data["max_storage_bytes"] = round(gb_value * (1024**3)) - return cls.model_validate(data) + return cls(storage_config=StorageConfig.from_disk(data)) except (tomllib.TOMLDecodeError, UnicodeDecodeError, ValidationError): logger.warning("Invalid config file, skipping...") return None