Pydantic <-> toml

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

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)

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

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

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