Upgrading to pydantic v2

gcp-airflow/composer.tf

@@ -14,7 +14,7 @@ resource "google_composer_environment" "zenml-airflow" {
     software_config {
       image_version = "composer-2-airflow-2"
       pypi_packages = {
-        pydantic = "~=1.9.2"
+        pydantic = "~=2.7.1"
       }
     }
 

pyproject.toml

@@ -66,15 +66,15 @@ exclude = [
 # pydantic = {version = "2.0.2"}
 # compatible with Core ZenML
 python = ">=3.8,<3.12"
-pydantic = { version = "<1.11,>=1.9.0" }
+pydantic = { version = "~2.7" }
 pyyaml = { version = ">=6.0.1" }
 click = { version = "^8.0.1,<8.1.4" }
 python-terraform = { version = "^0.10.1" }
 rich = { version = "^12.0.0" }
 analytics-python = { version = "^1.4.0" }
 
 # dev dependencies
-ruff = { version = "^0.1.7", optional = true }
+ruff = { version = ">=0.1.7", optional = true }
 pytest = { version = "^7.4.0", optional = true }
 mypy = { version = "^1.7.1", optional = true }
 darglint = { version = "^1.8.1", optional = true }

src/mlstacks/cli/cli.py

@@ -11,6 +11,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """CLI for mlstacks."""
+
 import random
 import shutil
 import string

src/mlstacks/models/component.py

@@ -12,9 +12,9 @@
 #  permissions and limitations under the License.
 """Component model."""
 
-from typing import Any, Dict, Optional
+from typing import Dict, Optional
 
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, field_validator, model_validator
 
 from mlstacks.constants import (
     INVALID_COMPONENT_FLAVOR_ERROR_MESSAGE,
@@ -67,8 +67,9 @@ class Component(BaseModel):
     component_flavor: ComponentFlavorEnum
     metadata: Optional[ComponentMetadata] = None
 
-    @validator("name")
-    def validate_name(cls, name: str) -> str:  # noqa
+    @field_validator("name")
+    @classmethod
+    def validate_name(cls, name: str) -> str:
         """Validate the name.
 
         Name must start with an alphanumeric character and can only contain
@@ -89,55 +90,35 @@ def validate_name(cls, name: str) -> str:  # noqa
             raise ValueError(INVALID_NAME_ERROR_MESSAGE)
         return name
 
-    @validator("component_type")
-    def validate_component_type(
-        cls,  # noqa
-        component_type: str,
-        values: Dict[str, Any],
-    ) -> str:
-        """Validate the component type.
+    @model_validator(mode="after")
+    def validate_component_type_and_flavor(self) -> "Component":
+        """Validate the component type and flavor.
 
         Artifact Store, Container Registry, Experiment Tracker, Orchestrator,
         MLOps Platform, and Model Deployer may be used with aws, gcp, and k3d
         providers. Step Operator may only be used with aws and gcp.
 
-        Args:
-            component_type: The component type.
-            values: The previously validated component specs.
+        Moreover, only certain flavors are allowed for a given
+        provider-component type combination. For more information, consult
+        the tables for your specified provider at the MLStacks documentation:
+        https://mlstacks.zenml.io/stacks/stack-specification.
 
         Returns:
-            The validated component type.
+            The validated component instance.
 
         Raises:
-            ValueError: If the component type is invalid.
+            ValueError: If the component type or flavor is invalid.
         """
-        if not is_valid_component_type(component_type, values["provider"]):
+        if not is_valid_component_type(
+            component_type=self.component_type, provider=self.provider
+        ):
             raise ValueError(INVALID_COMPONENT_TYPE_ERROR_MESSAGE)
-        return component_type
-
-    @validator("component_flavor")
-    def validate_component_flavor(
-        cls,  # noqa
-        component_flavor: str,
-        values: Dict[str, Any],
-    ) -> str:
-        """Validate the component flavor.
-
-        Only certain flavors are allowed for a given provider-component
-        type combination. For more information, consult the tables for
-        your specified provider at the MLStacks documentation:
-        https://mlstacks.zenml.io/stacks/stack-specification.
-
-        Args:
-            component_flavor: The component flavor.
-            values: The previously validated component specs.
 
-        Returns:
-            The validated component flavor.
-
-        Raises:
-            ValueError: If the component flavor is invalid.
-        """
-        if not is_valid_component_flavor(component_flavor, values):
+        if not is_valid_component_flavor(
+            component_flavor=self.component_flavor,
+            component_type=self.component_type,
+            provider=self.provider,
+        ):
             raise ValueError(INVALID_COMPONENT_FLAVOR_ERROR_MESSAGE)
-        return component_flavor
+
+        return self

src/mlstacks/models/stack.py

@@ -11,9 +11,10 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Stack model."""
+
 from typing import Dict, List, Optional
 
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, field_validator
 
 from mlstacks.constants import INVALID_NAME_ERROR_MESSAGE
 from mlstacks.enums import (
@@ -44,15 +45,16 @@ class Stack(BaseModel):
     spec_type: SpecTypeEnum = SpecTypeEnum.STACK
     name: str
     provider: ProviderEnum
-    default_region: Optional[str]
+    default_region: Optional[str] = None
     default_tags: Optional[Dict[str, str]] = None
-    deployment_method: Optional[
-        DeploymentMethodEnum
-    ] = DeploymentMethodEnum.KUBERNETES
+    deployment_method: Optional[DeploymentMethodEnum] = (
+        DeploymentMethodEnum.KUBERNETES
+    )
     components: List[Component] = []
 
-    @validator("name")
-    def validate_name(cls, name: str) -> str:  # noqa
+    @field_validator("name")
+    @classmethod
+    def validate_name(cls, name: str) -> str:
         """Validate the name.
 
         Name must start with an alphanumeric character and can only contain

src/mlstacks/utils/environment_utils.py

@@ -11,6 +11,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Environment utilities for mlstacks."""
+
 import os
 
 

src/mlstacks/utils/model_utils.py

@@ -13,7 +13,6 @@
 """Util functions for Pydantic models and validation."""
 
 import re
-from typing import Any, Dict
 
 from mlstacks.constants import ALLOWED_COMPONENT_TYPES, PERMITTED_NAME_REGEX
 
@@ -49,25 +48,26 @@ def is_valid_component_type(component_type: str, provider: str) -> bool:
 
 
 def is_valid_component_flavor(
-    component_flavor: str, specs: Dict[str, Any]
+    component_flavor: str,
+    component_type: str,
+    provider: str,
 ) -> bool:
     """Check if the component flavor is valid.
 
     Used for components.
 
     Args:
         component_flavor: The component flavor.
-        specs: The previously validated component specs.
+        component_type: The component type.
+        provider: The provider.
 
     Returns:
         True if the component flavor is valid, False otherwise.
     """
     try:
         is_valid = (
             component_flavor
-            in ALLOWED_COMPONENT_TYPES[specs["provider"]][
-                specs["component_type"]
-            ]
+            in ALLOWED_COMPONENT_TYPES[provider][component_type]
         )
     except KeyError:
         return False

src/mlstacks/utils/yaml_utils.py

@@ -11,6 +11,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Utility functions for loading YAML files into Python objects."""
+
 from pathlib import Path
 from typing import Any, Dict, Union
 

tests/unit/utils/test_zenml_utils.py

@@ -12,7 +12,6 @@
 #  permissions and limitations under the License.
 """Tests for utilities for mlstacks-ZenML interaction."""
 
-
 from mlstacks.models.component import Component
 from mlstacks.models.stack import Stack
 from mlstacks.utils.zenml_utils import has_valid_flavor_combinations