From 29eaa4ab8d2ddda38b4f2d93e50ff95cf42cee9a Mon Sep 17 00:00:00 2001 From: Antonyjin Date: Fri, 27 Sep 2024 22:57:39 +0900 Subject: [PATCH] Fix: Refactor classes and functions to align with Pydantic v2. - Updated code to explicitly specify optional keys where necessary. - Replaced direct `.get` calls with `data.get()` to handle new validation logic. - Migrated model configuration to use `model_config = ConfigDict(extra="forbid")` or `model_config = ConfigDict(extra="allow")` in place of Pydantic v1's configuration style. --- aleph_message/models/abstract.py | 5 +- aleph_message/models/execution/environment.py | 25 ++++----- aleph_message/models/execution/instance.py | 12 ++++- aleph_message/models/execution/program.py | 8 ++- aleph_message/models/execution/volume.py | 51 ++++++++++++------- 5 files changed, 62 insertions(+), 39 deletions(-) diff --git a/aleph_message/models/abstract.py b/aleph_message/models/abstract.py index f272dbd..2af6f8b 100644 --- a/aleph_message/models/abstract.py +++ b/aleph_message/models/abstract.py @@ -1,4 +1,4 @@ -from pydantic import BaseModel, Extra +from pydantic import BaseModel, ConfigDict def hashable(obj): @@ -24,5 +24,4 @@ class BaseContent(BaseModel): address: str time: float - class Config: - extra = Extra.forbid + model_config = ConfigDict(extra="forbid") diff --git a/aleph_message/models/execution/environment.py b/aleph_message/models/execution/environment.py index e0eae1d..7c6a475 100644 --- a/aleph_message/models/execution/environment.py +++ b/aleph_message/models/execution/environment.py @@ -3,7 +3,7 @@ from enum import Enum from typing import List, Literal, Optional, Union -from pydantic import Extra, Field, validator +from pydantic import ConfigDict, Field, field_validator from ...utils import Mebibytes from ..abstract import HashableModel @@ -13,8 +13,7 @@ class Subscription(HashableModel): """A subscription is used to trigger a program in response to a FunctionTrigger.""" - class Config: - extra = Extra.allow + model_config = ConfigDict(extra="allow") class FunctionTriggers(HashableModel): @@ -29,8 +28,7 @@ class FunctionTriggers(HashableModel): description="Persist the execution of the program instead of running it on demand.", ) - class Config: - extra = Extra.forbid + model_config = ConfigDict(extra="forbid") class NetworkProtocol(str, Enum): @@ -85,8 +83,7 @@ class CpuProperties(HashableModel): description="CPU features required by the virtual machine. Examples: 'sev', 'sev_es', 'sev_snp'.", ) - class Config: - extra = Extra.forbid + model_config = ConfigDict(extra="forbid") class HypervisorType(str, Enum): @@ -132,8 +129,7 @@ class TrustedExecutionEnvironment(HashableModel): description="Policy of the TEE. Default value is 0x01 for SEV without debugging.", ) - class Config: - extra = Extra.allow + model_config = ConfigDict(extra="allow") class InstanceEnvironment(HashableModel): @@ -150,9 +146,9 @@ class InstanceEnvironment(HashableModel): reproducible: bool = False shared_cache: bool = False - @validator("trusted_execution", pre=True) + @field_validator("trusted_execution", mode="before") def check_hypervisor(cls, v, values): - if v and values.get("hypervisor") != HypervisorType.qemu: + if v and values.data.get("hypervisor") != HypervisorType.qemu: raise ValueError("Trusted Execution Environment is only supported for QEmu") return v @@ -166,8 +162,7 @@ class NodeRequirements(HashableModel): default=None, description="Hash of the compute resource node that must be used" ) - class Config: - extra = Extra.forbid + model_config = ConfigDict(extra="forbid") class HostRequirements(HashableModel): @@ -178,6 +173,4 @@ class HostRequirements(HashableModel): default=None, description="Required Compute Resource Node properties" ) - class Config: - # Allow users to add custom requirements - extra = Extra.allow + model_config = ConfigDict(extra="allow") diff --git a/aleph_message/models/execution/instance.py b/aleph_message/models/execution/instance.py index ebb8d48..4c54b25 100644 --- a/aleph_message/models/execution/instance.py +++ b/aleph_message/models/execution/instance.py @@ -1,6 +1,7 @@ from __future__ import annotations +from typing import Optional, List -from pydantic import Field +from pydantic import Field, field_validator from aleph_message.models.abstract import HashableModel @@ -21,11 +22,20 @@ class RootfsVolume(HashableModel): persistence: VolumePersistence # Use the same size constraint as persistent volumes for now size_mib: PersistentVolumeSizeMib + forgotten_by: Optional[List[str]] = None + + @field_validator('size_mib', mode="before") + def convert_size_mib(cls, v): + if isinstance(v, int): + return PersistentVolumeSizeMib(persistent_volume_size=v) + return v class InstanceContent(BaseExecutableContent): """Message content for scheduling a VM instance on the network.""" + metadata: Optional[dict] = None + payment: Optional[dict] = None environment: InstanceEnvironment = Field( description="Properties of the instance execution environment" ) diff --git a/aleph_message/models/execution/program.py b/aleph_message/models/execution/program.py index 8afb6d9..ee3237c 100644 --- a/aleph_message/models/execution/program.py +++ b/aleph_message/models/execution/program.py @@ -43,8 +43,8 @@ class DataContent(HashableModel): encoding: Encoding mount: str - ref: ItemHash - use_latest: bool = False + ref: Optional[ItemHash] = None + use_latest: Optional[bool] = False class Export(HashableModel): @@ -69,3 +69,7 @@ class ProgramContent(BaseExecutableContent): default=None, description="Data to export after computation" ) on: FunctionTriggers = Field(description="Signals that trigger an execution") + + metadata: Optional[dict] = None + authorized_keys: Optional[List[str]] = None + payment: Optional[dict] = None diff --git a/aleph_message/models/execution/volume.py b/aleph_message/models/execution/volume.py index 6102fba..adbfe3c 100644 --- a/aleph_message/models/execution/volume.py +++ b/aleph_message/models/execution/volume.py @@ -4,7 +4,7 @@ from enum import Enum from typing import Literal, Optional, Union -from pydantic import ConstrainedInt, Extra +from pydantic import Field, ConfigDict, BaseModel, field_validator from ...utils import Gigabytes, gigabyte_to_mebibyte from ..abstract import HashableModel @@ -18,27 +18,35 @@ class AbstractVolume(HashableModel, ABC): @abstractmethod def is_read_only(self): ... - class Config: - extra = Extra.forbid + model_config = ConfigDict(extra="forbid") class ImmutableVolume(AbstractVolume): - ref: ItemHash + ref: ItemHash = None use_latest: bool = True def is_read_only(self): return True -class EphemeralVolumeSize(ConstrainedInt): - gt = 0 - le = 1000 # Limit to 1 GiB - strict = True +class EphemeralVolumeSize(BaseModel): + ephemeral_volume_size: int = Field(gt=0, + le=1000, #Limit to 1GiB + strict=True) + + def __hash__(self): + return hash(self.ephemeral_volume_size) class EphemeralVolume(AbstractVolume): ephemeral: Literal[True] = True - size_mib: EphemeralVolumeSize + size_mib: EphemeralVolumeSize = 0 + + @field_validator('size_mib', mode="before") + def convert_size_mib(cls, v): + if isinstance(v, int): + return EphemeralVolumeSize(ephemeral_volume_size=v) + return v def is_read_only(self): return False @@ -58,17 +66,26 @@ class VolumePersistence(str, Enum): store = "store" -class PersistentVolumeSizeMib(ConstrainedInt): - gt = 0 - le = gigabyte_to_mebibyte(Gigabytes(100)) - strict = True # Limit to 100 GiB +class PersistentVolumeSizeMib(BaseModel): + persistent_volume_size: int = Field(gt=0, + le=gigabyte_to_mebibyte(Gigabytes(100)), #Limit to 1GiB + strict=True) + + def __hash__(self): + return hash(self.persistent_volume_size) class PersistentVolume(AbstractVolume): - parent: Optional[ParentVolume] - persistence: VolumePersistence - name: str - size_mib: PersistentVolumeSizeMib + parent: Optional[ParentVolume] = None + persistence: VolumePersistence = None + name: Optional[str] = None + size_mib: PersistentVolumeSizeMib = 0 + + @field_validator('size_mib', mode="before") + def convert_size_mib(cls, v): + if isinstance(v, int): + return PersistentVolumeSizeMib(persistent_volume_size=v) + return v def is_read_only(self): return False