zenml-nightly 0.58.2.dev20240614__py3-none-any.whl → 0.58.2.dev20240622__py3-none-any.whl

zenml/orchestrators/step_launcher.py

@@ -17,7 +17,8 @@ import os
 import time
 from contextlib import nullcontext
 from datetime import datetime
-from typing import TYPE_CHECKING, Dict, Optional, Tuple
+from functools import partial
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple
 
 from zenml.client import Client
 from zenml.config.step_configurations import Step
@@ -247,10 +248,28 @@ class StepLauncher:
                     while retries < max_retries:
                         last_retry = retries == max_retries - 1
                         try:
+                            # here pass a forced save_to_file callable to be
+                            # used as a dump function to use before starting
+                            # the external jobs in step operators
+                            if isinstance(
+                                logs_context,
+                                step_logging.StepLogsStorageContext,
+                            ):
+                                force_write_logs = partial(
+                                    logs_context.storage.save_to_file,
+                                    force=True,
+                                )
+                            else:
+
+                                def _bypass() -> None:
+                                    return None
+
+                                force_write_logs = _bypass
                             self._run_step(
                                 pipeline_run=pipeline_run,
                                 step_run=step_run_response,
                                 last_retry=last_retry,
+                                force_write_logs=force_write_logs,
                             )
                             logger.info(
                                 f"Step `{self._step_name}` completed successfully."
@@ -324,9 +343,11 @@ class StepLauncher:
             user=client.active_user.id,
             workspace=client.active_workspace.id,
             deployment=self._deployment.id,
-            pipeline=self._deployment.pipeline.id
-            if self._deployment.pipeline
-            else None,
+            pipeline=(
+                self._deployment.pipeline.id
+                if self._deployment.pipeline
+                else None
+            ),
             status=ExecutionStatus.RUNNING,
             orchestrator_environment=get_run_environment_dict(),
             start_time=datetime.utcnow(),
@@ -415,6 +436,7 @@ class StepLauncher:
         self,
         pipeline_run: PipelineRunResponse,
         step_run: StepRunResponse,
+        force_write_logs: Callable[..., Any],
         last_retry: bool = True,
     ) -> None:
         """Runs the current step.
@@ -422,6 +444,7 @@ class StepLauncher:
         Args:
             pipeline_run: The model of the current pipeline run.
             step_run: The model of the current step run.
+            force_write_logs: The context for the step logs.
             last_retry: Whether this is the last retry of the step.
         """
         # Prepare step run information.
@@ -432,6 +455,7 @@ class StepLauncher:
             pipeline_step_name=self._step_name,
             run_id=pipeline_run.id,
             step_run_id=step_run.id,
+            force_write_logs=force_write_logs,
         )
 
         output_artifact_uris = output_utils.prepare_output_artifact_uris(

zenml/orchestrators/step_runner.py

@@ -28,8 +28,6 @@ from typing import (
 )
 from uuid import UUID
 
-from pydantic.typing import get_origin, is_union
-
 from zenml.artifacts.unmaterialized_artifact import UnmaterializedArtifact
 from zenml.artifacts.utils import save_artifact
 from zenml.config.step_configurations import StepConfiguration
@@ -63,6 +61,7 @@ from zenml.steps.utils import (
     resolve_type_annotation,
 )
 from zenml.utils import materializer_utils, source_utils
+from zenml.utils.typing_utils import get_origin, is_union
 
 if TYPE_CHECKING:
     from zenml.config.source import Source
@@ -396,7 +395,7 @@ class StepRunner:
 
             # Parse the parameters
             if issubclass(arg_type, BaseParameters):
-                step_params = arg_type.parse_obj(
+                step_params = arg_type.model_validate(
                     self.configuration.parameters[arg]
                 )
                 function_params[arg] = step_params
@@ -440,7 +439,7 @@ class StepRunner:
         # Skip materialization for `UnmaterializedArtifact`.
         if data_type == UnmaterializedArtifact:
             return UnmaterializedArtifact(
-                **artifact.get_hydrated_version().dict()
+                **artifact.get_hydrated_version().model_dump()
             )
 
         if data_type is Any or is_union(get_origin(data_type)):
@@ -513,8 +512,6 @@ class StepRunner:
                 f"(return values: {return_values})."
             )
 
-        from pydantic.typing import get_origin, is_union
-
         from zenml.steps.utils import get_args
 
         validated_outputs: Dict[str, Any] = {}

zenml/orchestrators/utils.py

@@ -225,7 +225,7 @@ def _link_cached_artifacts_to_model(
         if output_name_ in output_annotations:
             annotation = output_annotations.get(output_name_, None)
             if annotation and annotation.artifact_config is not None:
-                artifact_config_ = annotation.artifact_config.copy()
+                artifact_config_ = annotation.artifact_config.model_copy()
         # no artifact config found or artifact was produced by `save_artifact`
         # inside the step body, so was never in annotations
         if artifact_config_ is None:

zenml/plugins/base_plugin_flavor.py

@@ -16,7 +16,7 @@
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, ClassVar, Type
 
-from pydantic import BaseModel, Extra
+from pydantic import BaseModel, ConfigDict
 
 from zenml.config.global_config import GlobalConfiguration
 from zenml.enums import PluginSubType, PluginType
@@ -29,16 +29,12 @@ if TYPE_CHECKING:
 class BasePluginConfig(BaseModel, ABC):
     """Allows configuring of Event Source and Filter configuration."""
 
-    class Config:
-        """Pydantic configuration class."""
-
-        # public attributes are immutable
-        allow_mutation = True
-        # all attributes with leading underscore are private and therefore
-        # are mutable and not included in serialization
-        underscore_attrs_are_private = True
+    model_config = ConfigDict(
+        # public attributes are mutable
+        frozen=False,
         # ignore extra attributes during model initialization
-        extra = Extra.ignore
+        extra="ignore",
+    )
 
 
 class BasePlugin(ABC):

zenml/plugins/plugin_flavor_registry.py

@@ -16,7 +16,7 @@
 import math
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Type
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 from zenml.enums import PluginSubType, PluginType
 from zenml.integrations.registry import integration_registry
@@ -35,12 +35,8 @@ class RegistryEntry(BaseModel):
     """Registry Entry Class for the Plugin Registry."""
 
     flavor_class: Type[BasePluginFlavor]
-    plugin_instance: Optional[BasePlugin]
-
-    class Config:
-        """Pydantic configuration class."""
-
-        arbitrary_types_allowed = True
+    plugin_instance: Optional[BasePlugin] = None
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
 
 class PluginFlavorRegistry:

zenml/secret/base_secret.py

@@ -15,7 +15,7 @@
 
 from typing import Any, Dict, List
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 
 class BaseSecretSchema(BaseModel):
@@ -31,7 +31,7 @@ class BaseSecretSchema(BaseModel):
         Returns:
             A list of all attribute names that are part of the schema.
         """
-        return list(cls.__fields__.keys())
+        return list(cls.model_fields.keys())
 
     def get_values(self) -> Dict[str, Any]:
         """Get all values of the secret schema.
@@ -39,12 +39,11 @@ class BaseSecretSchema(BaseModel):
         Returns:
             A dictionary of all attribute names and their corresponding values.
         """
-        return self.dict(exclude_none=True)
-
-    class Config:
-        """Pydantic configuration class."""
+        return self.model_dump(exclude_none=True)
 
+    model_config = ConfigDict(
         # validate attribute assignments
-        validate_assignment = True
+        validate_assignment=True,
         # report extra attributes as validation failures
-        extra = "forbid"
+        extra="forbid",
+    )

zenml/service_connectors/docker_service_connector.py

@@ -23,7 +23,7 @@ from typing import Any, List, Optional
 
 from docker.client import DockerClient
 from docker.errors import DockerException
-from pydantic import Field, SecretStr
+from pydantic import Field
 
 from zenml.constants import DOCKER_REGISTRY_RESOURCE_TYPE
 from zenml.exceptions import AuthorizationException
@@ -39,6 +39,7 @@ from zenml.service_connectors.service_connector import (
 )
 from zenml.utils import docker_utils
 from zenml.utils.enum_utils import StrEnum
+from zenml.utils.secret_utils import PlainSerializedSecretStr
 
 logger = get_logger(__name__)
 
@@ -46,10 +47,10 @@ logger = get_logger(__name__)
 class DockerCredentials(AuthenticationConfig):
     """Docker client authentication credentials."""
 
-    username: SecretStr = Field(
+    username: PlainSerializedSecretStr = Field(
         title="Username",
     )
-    password: SecretStr = Field(
+    password: PlainSerializedSecretStr = Field(
         title="Password",
     )
 

zenml/service_connectors/service_connector.py

@@ -34,7 +34,7 @@ from pydantic import (
     SecretStr,
     ValidationError,
 )
-from pydantic.main import ModelMetaclass
+from pydantic._internal._model_construction import ModelMetaclass
 
 from zenml.client import Client
 from zenml.constants import (
@@ -71,7 +71,7 @@ class AuthenticationConfig(BaseModel):
         """
         return {
             k: v
-            for k, v in self.dict(exclude_none=True).items()
+            for k, v in self.model_dump(exclude_none=True).items()
             if isinstance(v, SecretStr)
         }
 
@@ -84,7 +84,7 @@ class AuthenticationConfig(BaseModel):
         """
         return {
             k: v
-            for k, v in self.dict(exclude_none=True).items()
+            for k, v in self.model_dump(exclude_none=True).items()
             if not isinstance(v, SecretStr)
         }
 
@@ -95,7 +95,7 @@ class AuthenticationConfig(BaseModel):
         Returns:
             A dictionary of all values in the configuration.
         """
-        return self.dict(exclude_none=True)
+        return self.model_dump(exclude_none=True)
 
 
 class ServiceConnectorMeta(ModelMetaclass):
@@ -487,7 +487,7 @@ class ServiceConnector(BaseModel, metaclass=ServiceConnectorMeta):
         ):
             return self
 
-        copy = self.copy()
+        copy = self.model_copy()
         copy.resource_type = resource_type
         copy.resource_id = resource_id
         return copy
@@ -1401,10 +1401,3 @@ class ServiceConnector(BaseModel, metaclass=ServiceConnectorMeta):
         )
 
         return connector_client
-
-    class Config:
-        """Connector configuration."""
-
-        # all attributes with leading underscore are private and therefore
-        # are mutable and not included in serialization
-        underscore_attrs_are_private = True

zenml/service_connectors/service_connector_registry.py

@@ -22,7 +22,6 @@ if TYPE_CHECKING:
     from zenml.models import (
         ServiceConnectorRequest,
         ServiceConnectorResponse,
-        ServiceConnectorUpdate,
     )
     from zenml.service_connectors.service_connector import ServiceConnector
 logger = get_logger(__name__)
@@ -90,7 +89,7 @@ class ServiceConnectorRegistry:
                 f"Please make sure the corresponding packages and/or ZenML "
                 f"integration are installed and try again."
             )
-        return self.service_connector_types[connector_type].copy()
+        return self.service_connector_types[connector_type].model_copy()
 
     def __getitem__(self, key: str) -> ServiceConnectorTypeModel:
         """Get a service connector type by its connector type identifier.
@@ -155,7 +154,7 @@ class ServiceConnectorRegistry:
                     or auth_method in service_connector_type.auth_method_dict
                 )
             ):
-                matches.append(service_connector_type.copy())
+                matches.append(service_connector_type.model_copy())
 
         return matches
 
@@ -164,7 +163,6 @@ class ServiceConnectorRegistry:
         model: Union[
             "ServiceConnectorRequest",
             "ServiceConnectorResponse",
-            "ServiceConnectorUpdate",
         ],
     ) -> "ServiceConnector":
         """Validate a service connector model and create an instance from it.

zenml/services/container/container_service.py

@@ -289,7 +289,7 @@ class ContainerService(BaseService):
         assert self.status.log_file is not None
 
         with open(self.status.config_file, "w") as f:
-            f.write(self.json(indent=4))
+            f.write(self.model_dump_json(indent=4))
         pathlib.Path(self.status.log_file).touch()
 
         command = [

zenml/services/container/container_service_endpoint.py

@@ -80,7 +80,7 @@ class ContainerServiceEndpoint(BaseServiceEndpoint):
     )
     monitor: Optional[
         Union[HTTPEndpointHealthMonitor, TCPEndpointHealthMonitor]
-    ] = Field(..., discriminator="type")
+    ] = Field(..., discriminator="type", union_mode="left_to_right")
 
     def _lookup_free_port(self) -> int:
         """Search for a free TCP port for the service endpoint.

zenml/services/local/local_service.py

@@ -331,7 +331,7 @@ class LocalDaemonService(BaseService):
         assert self.status.pid_file is not None
 
         with open(self.status.config_file, "w") as f:
-            f.write(self.json(indent=4))
+            f.write(self.model_dump_json(indent=4))
 
         # delete the previous PID file, in case a previous daemon process
         # crashed and left a stale PID file

zenml/services/local/local_service_endpoint.py

@@ -83,7 +83,7 @@ class LocalDaemonServiceEndpoint(BaseServiceEndpoint):
     )
     monitor: Optional[
         Union[HTTPEndpointHealthMonitor, TCPEndpointHealthMonitor]
-    ] = Field(..., discriminator="type")
+    ] = Field(..., discriminator="type", union_mode="left_to_right")
 
     def _lookup_free_port(self) -> int:
         """Search for a free TCP port for the service endpoint.

zenml/services/service.py

@@ -31,6 +31,8 @@ from typing import (
 )
 from uuid import UUID
 
+from pydantic import ConfigDict
+
 from zenml.console import console
 from zenml.logger import get_logger
 from zenml.services.service_endpoint import BaseServiceEndpoint
@@ -117,6 +119,14 @@ class ServiceConfig(BaseTypedModel):
     model_version: str = ""
     service_name: str = ""
 
+    # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
+    #  fields defined under base models. If not handled, this raises a warning.
+    #  It is possible to suppress this warning message with the following
+    #  configuration, however the ultimate solution is to rename these fields.
+    #  Even though they do not cause any problems right now, if we are not
+    #  careful we might overwrite some fields protected by pydantic.
+    model_config = ConfigDict(protected_namespaces=())
+
     def __init__(self, **data: Any):
         """Initialize the service configuration.
 
@@ -142,7 +152,7 @@ class ServiceConfig(BaseTypedModel):
             a dictionary of service labels.
         """
         labels = {}
-        for k, v in self.dict().items():
+        for k, v in self.model_dump().items():
             label = f"zenml_{k}".upper()
             labels[label] = str(v)
         return labels
@@ -172,7 +182,7 @@ class BaseService(BaseTypedModel):
     config: ServiceConfig
     status: ServiceStatus
     # TODO [ENG-703]: allow multiple endpoints per service
-    endpoint: Optional[BaseServiceEndpoint]
+    endpoint: Optional[BaseServiceEndpoint] = None
 
     def __init__(
         self,
@@ -209,7 +219,7 @@ class BaseService(BaseTypedModel):
             admin_state=model.admin_state,
             config=model.config,
             status=model.status,
-            service_type=model.service_type.dict(),
+            service_type=model.service_type.model_dump(),
             endpoint=model.endpoint,
         )
 
@@ -512,14 +522,10 @@ class BaseService(BaseTypedModel):
         """
         return self.__repr__()
 
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # validate attribute assignments
-        validate_assignment = True
-        # all attributes with leading underscore are private and therefore
-        # are mutable and not included in serialization
-        underscore_attrs_are_private = True
+        validate_assignment=True,
+    )
 
 
 class BaseDeploymentService(BaseService):

zenml/services/service_type.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Implementation of a ZenML ServiceType class."""
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 
 class ServiceType(BaseModel):
@@ -33,8 +33,7 @@ class ServiceType(BaseModel):
     description: str = ""
     logo_url: str = ""
 
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # make the service type immutable and hashable
-        frozen = True
+        frozen=True
+    )

zenml/services/terraform/terraform_service.py

@@ -176,7 +176,7 @@ class TerraformService(BaseService):
         assert self.status.config_file is not None
 
         with open(self.status.config_file, "w") as f:
-            f.write(self.json(indent=4))
+            f.write(self.model_dump_json(indent=4))
 
     def _write_vars_to_file(self, vars: Dict[str, Any]) -> None:
         """Write variables to the variables file.

zenml/stack/flavor.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Base ZenML Flavor implementation."""
 
-import json
 from abc import abstractmethod
 from typing import Any, Dict, Optional, Type, cast
 
@@ -102,10 +101,7 @@ class Flavor:
         Returns:
             The config schema.
         """
-        config_schema: Dict[str, Any] = json.loads(
-            self.config_class.schema_json()
-        )
-        return config_schema
+        return self.config_class.model_json_schema()
 
     @property
     def service_connector_requirements(
@@ -209,7 +205,7 @@ class Flavor:
             is_latest = True
 
         if is_latest:
-            base = "https://docs.zenml.io/"
+            base = "https://docs.zenml.io"
         else:
             base = f"https://zenml-io.gitbook.io/zenml-legacy-documentation/v/{__version__}"
         return f"{base}/stack-components/{component_type}/{name}"

zenml/stack/flavor_registry.py

@@ -122,8 +122,8 @@ class FlavorRegistry:
                 if len(existing_flavor) == 0:
                     store.create_flavor(flavor_request_model)
                 else:
-                    flavor_update_model = FlavorUpdate.parse_obj(
-                        flavor_request_model
+                    flavor_update_model = FlavorUpdate.model_validate(
+                        dict(flavor_request_model)
                     )
                     store.update_flavor(
                         existing_flavor[0].id, flavor_update_model
@@ -154,8 +154,8 @@ class FlavorRegistry:
                         if len(existing_flavor) == 0:
                             store.create_flavor(flavor_request_model)
                         else:
-                            flavor_update_model = FlavorUpdate.parse_obj(
-                                flavor_request_model
+                            flavor_update_model = FlavorUpdate.model_validate(
+                                dict(flavor_request_model)
                             )
                             store.update_flavor(
                                 existing_flavor[0].id, flavor_update_model

zenml/stack/stack.py

@@ -15,6 +15,7 @@
 
 import functools
 import itertools
+import json
 import os
 from datetime import datetime
 from typing import (
@@ -481,7 +482,9 @@ class Stack:
             A dictionary containing the stack components.
         """
         component_dict = {
-            component_type.value: component.config.json(sort_keys=True)
+            component_type.value: json.dumps(
+                component.config.model_dump(mode="json"), sort_keys=True
+            )
             for component_type, component in self.components.items()
         }
         component_dict.update({"name": self.name})