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

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

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

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

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

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