diff --git a/src/exo/shared/types/commands.py b/src/exo/shared/types/commands.py index a6c988a5..c5041e07 100644 --- a/src/exo/shared/types/commands.py +++ b/src/exo/shared/types/commands.py @@ -7,6 +7,8 @@ from exo.api.types import ( 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.memory import Memory +from exo.shared.types.storage import StoragePolicy from exo.shared.types.text_generation import TextGenerationTaskParams from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta from exo.shared.types.worker.shards import Sharding, ShardMetadata @@ -92,6 +94,12 @@ class DeleteCustomModelCard(BaseCommand): DownloadCommand = StartDownload | DeleteDownload | CancelDownload +class SetStorageConfig(BaseCommand): + target_node_id: NodeId + max_storage: Memory | None + storage_policy: StoragePolicy + + Command = ( TestCommand | RequestEventLog @@ -106,6 +114,7 @@ Command = ( | SendInputChunk | AddCustomModelCard | DeleteCustomModelCard + | SetStorageConfig )