zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/services/container/entrypoint.py

@@ -19,6 +19,7 @@ that are implemented as locally running docker containers.
 
 import os
 import sys
+from typing import cast
 
 import click
 
@@ -50,7 +51,7 @@ def launch_service(service_config_file: str) -> None:
     # with messages before daemonization is complete
     from zenml.integrations.registry import integration_registry
     from zenml.logger import get_logger
-    from zenml.services import ContainerService, ServiceRegistry
+    from zenml.services import ContainerService
 
     logger = get_logger(__name__)
 
@@ -63,7 +64,7 @@ def launch_service(service_config_file: str) -> None:
     logger.debug(
         "Running containerized service with configuration:\n %s", config
     )
-    service = ServiceRegistry().load_service_from_json(config)
+    service = cast("ContainerService", ContainerService.from_json(config))
     if not isinstance(service, ContainerService):
         raise TypeError(
             f"Expected service type ContainerService but got "

zenml/services/local/local_daemon_entrypoint.py

@@ -18,18 +18,19 @@ that are implemented as locally running daemon processes.
 """
 
 import os
+from typing import cast
 
 import click
 
 from zenml.utils.daemon import daemonize
 
-# Try to import the LocalZenServer here because it needs to be registered in the
+# Try to import the DaemonZenServer here because it needs to be registered in the
 # service registry early on in order to be available for use in other modules.
-# If the LocalZenServer dependencies aren't installed, there is no need to register
+# If the DaemonZenServer dependencies aren't installed, there is no need to register
 # it anywhere so we simply pass.
 try:
-    from zenml.zen_server.deploy.local.local_zen_server import (  # noqa
-        LocalZenServer,
+    from zenml.zen_server.deploy.daemon.daemon_zen_server import (  # noqa
+        DaemonZenServer,
     )
 except ImportError:
     pass
@@ -68,7 +69,7 @@ def run(
         # with messages before daemonization is complete
         from zenml.integrations.registry import integration_registry
         from zenml.logger import get_logger
-        from zenml.services import LocalDaemonService, ServiceRegistry
+        from zenml.services import LocalDaemonService
 
         logger = get_logger(__name__)
 
@@ -81,7 +82,9 @@ def run(
         integration_registry.activate_integrations()
 
         logger.debug("Running service daemon with configuration:\n %s", config)
-        service = ServiceRegistry().load_service_from_json(config)
+        service = cast(
+            "LocalDaemonService", LocalDaemonService.from_json(config)
+        )
         if not isinstance(service, LocalDaemonService):
             raise TypeError(
                 f"Expected service type LocalDaemonService but got "

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

@@ -13,23 +13,88 @@
 #  permissions and limitations under the License.
 """Implementation of the ZenML Service class."""
 
+import json
 import time
 from abc import abstractmethod
-from typing import Any, ClassVar, Dict, Generator, Optional, Tuple, Type, cast
-from uuid import UUID, uuid4
-
-from pydantic import Field
+from functools import wraps
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    ClassVar,
+    Dict,
+    Generator,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+)
+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
-from zenml.services.service_registry import ServiceRegistry
+from zenml.services.service_monitor import HTTPEndpointHealthMonitor
 from zenml.services.service_status import ServiceState, ServiceStatus
 from zenml.services.service_type import ServiceType
-from zenml.utils.typed_model import BaseTypedModel, BaseTypedModelMeta
+from zenml.utils import source_utils
+from zenml.utils.typed_model import BaseTypedModel
 
 logger = get_logger(__name__)
 
+T = TypeVar("T", bound=Callable[..., Any])
+
+if TYPE_CHECKING:
+    from zenml.models.v2.core.service import ServiceResponse
+
+ZENM_ENDPOINT_PREFIX = "zenml-"
+
+
+def update_service_status(
+    pre_status: Optional[ServiceState] = None,
+    post_status: Optional[ServiceState] = None,
+    error_status: ServiceState = ServiceState.ERROR,
+) -> Callable[[T], T]:
+    """A decorator to update the service status before and after a method call.
+
+    This decorator is used to wrap service methods and update the service status
+    before and after the method call. If the method raises an exception, the
+    service status is updated to reflect the error state.
+
+    Args:
+        pre_status: the status to update before the method call.
+        post_status: the status to update after the method call.
+        error_status: the status to update if the method raises an exception.
+
+    Returns:
+        Callable[..., Any]: The wrapped method with exception handling.
+    """
+
+    def decorator(func: T) -> T:
+        @wraps(func)
+        def wrapper(self: "BaseService", *args: Any, **kwargs: Any) -> Any:
+            if pre_status:
+                self.status.update_state(pre_status, "")
+            try:
+                logger.debug(f"Calling {func.__name__} method...")
+                result = func(self, *args, **kwargs)
+                logger.debug(f"{func.__name__} method executed successfully.")
+                if post_status:
+                    self.status.update_state(post_status, "")
+                return result
+            except Exception as e:
+                logger.error(
+                    f"Error occurred in {func.__name__} method: {str(e)}"
+                )
+                self.status.update_state(error_status, str(e))
+                raise
+
+        return wrapper  # type: ignore
+
+    return decorator
+
 
 class ServiceConfig(BaseTypedModel):
     """Generic service configuration.
@@ -50,107 +115,50 @@ class ServiceConfig(BaseTypedModel):
     description: str = ""
     pipeline_name: str = ""
     pipeline_step_name: str = ""
-    run_name: str = ""
-
-
-class BaseServiceMeta(BaseTypedModelMeta):
-    """Metaclass responsible for registering different BaseService subclasses.
+    model_name: str = ""
+    model_version: str = ""
+    service_name: str = ""
 
-    This metaclass has two main responsibilities:
-    1. register all BaseService types in the service registry. This is relevant
-    when services are deserialized and instantiated from their JSON or dict
-    representation, because their type needs to be known beforehand.
-    2. ensuring BaseService instance uniqueness by enforcing that no two
-    service instances have the same UUID value. Implementing this at the
-    constructor level guarantees that deserializing a service instance from
-    a JSON representation multiple times always returns the same service object.
-    """
+    # 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 __new__(
-        mcs, name: str, bases: Tuple[Type[Any], ...], dct: Dict[str, Any]
-    ) -> "BaseServiceMeta":
-        """Creates a BaseService class and registers it in the `ServiceRegistry`.
+    def __init__(self, **data: Any):
+        """Initialize the service configuration.
 
         Args:
-            name: name of the class.
-            bases: tuple of base classes.
-            dct: dictionary of class attributes.
-
-        Returns:
-            the created BaseServiceMeta class.
+            **data: keyword arguments.
 
         Raises:
-            TypeError: if the 'service_type' reserved attribute name is used.
+            ValueError: if neither 'name' nor 'model_name' is set.
         """
-        service_type = dct.get("SERVICE_TYPE", None)
-
-        # register only classes of concrete service implementations
-        if service_type:
-            # add the service type class attribute to the class as a regular
-            # immutable attribute to include it in the JSON representation
-            if "service_type" in dct:
-                raise TypeError(
-                    "`service_type` is a reserved attribute name for BaseService "
-                    "subclasses"
-                )
-            dct.setdefault("__annotations__", dict())[
-                "service_type"
-            ] = ServiceType
-            dct["service_type"] = Field(service_type, allow_mutation=False)
-
-        cls = cast(Type["BaseService"], super().__new__(mcs, name, bases, dct))
-
-        # register only classes of concrete service implementations
-        if service_type:
-            # register the service type in the service registry
-            ServiceRegistry().register_service_type(cls)
-        return cls
+        super().__init__(**data)
+        if self.name or self.model_name:
+            self.service_name = data.get(
+                "service_name",
+                f"{ZENM_ENDPOINT_PREFIX}{self.name or self.model_name}",
+            )
+        else:
+            raise ValueError("Either 'name' or 'model_name' must be set.")
 
-    def __call__(cls, *args: Any, **kwargs: Any) -> "BaseServiceMeta":
-        """Validate the creation of a service.
-
-        Args:
-            *args: positional arguments.
-            **kwargs: keyword arguments.
+    def get_service_labels(self) -> Dict[str, str]:
+        """Get the service labels.
 
         Returns:
-            the created BaseServiceMeta class.
-
-        Raises:
-            AttributeError: if the service UUID is untyped.
-            ValueError: if the service UUID is not a UUID type.
+            a dictionary of service labels.
         """
-        if not getattr(cls, "SERVICE_TYPE", None):
-            raise AttributeError(
-                f"Untyped service instances are not allowed. Please set the "
-                f"SERVICE_TYPE class attribute for {cls}."
-            )
-        uuid = kwargs.get("uuid", None)
-        if uuid:
-            if isinstance(uuid, str):
-                uuid = UUID(uuid)
-            if not isinstance(uuid, UUID):
-                raise ValueError(
-                    f"The `uuid` argument for {cls} must be a UUID instance or a "
-                    f"string representation of a UUID."
-                )
-
-            # if a service instance with the same UUID is already registered,
-            # return the existing instance rather than the newly created one
-            existing_service = ServiceRegistry().get_service(uuid)
-            if existing_service:
-                logger.debug(
-                    f"Reusing existing service '{existing_service}' "
-                    f"instead of creating a new service with the same UUID."
-                )
-                return cast("BaseServiceMeta", existing_service)
+        labels = {}
+        for k, v in self.model_dump().items():
+            label = f"zenml_{k}".upper()
+            labels[label] = str(v)
+        return labels
 
-        svc = cast("BaseService", super().__call__(*args, **kwargs))
-        ServiceRegistry().register_service(svc)
-        return cast("BaseServiceMeta", svc)
 
-
-class BaseService(BaseTypedModel, metaclass=BaseServiceMeta):
+class BaseService(BaseTypedModel):
     """Base service class.
 
     This class implements generic functionality concerning the life-cycle
@@ -169,12 +177,12 @@ class BaseService(BaseTypedModel, metaclass=BaseServiceMeta):
 
     SERVICE_TYPE: ClassVar[ServiceType]
 
-    uuid: UUID = Field(default_factory=uuid4, allow_mutation=False)
+    uuid: UUID
     admin_state: ServiceState = ServiceState.INACTIVE
     config: ServiceConfig
     status: ServiceStatus
     # TODO [ENG-703]: allow multiple endpoints per service
-    endpoint: Optional[BaseServiceEndpoint]
+    endpoint: Optional[BaseServiceEndpoint] = None
 
     def __init__(
         self,
@@ -188,6 +196,49 @@ class BaseService(BaseTypedModel, metaclass=BaseServiceMeta):
         super().__init__(**attrs)
         self.config.name = self.config.name or self.__class__.__name__
 
+    @classmethod
+    def from_model(cls, model: "ServiceResponse") -> "BaseService":
+        """Loads a service from a model.
+
+        Args:
+            model: The ServiceResponse to load from.
+
+        Returns:
+            The loaded service object.
+
+        Raises:
+            ValueError: if the service source is not found in the model.
+        """
+        if not model.service_source:
+            raise ValueError("Service source not found in the model.")
+        class_: Type[BaseService] = source_utils.load_and_validate_class(
+            source=model.service_source, expected_class=BaseService
+        )
+        return class_(
+            uuid=model.id,
+            admin_state=model.admin_state,
+            config=model.config,
+            status=model.status,
+            service_type=model.service_type.model_dump(),
+            endpoint=model.endpoint,
+        )
+
+    @classmethod
+    def from_json(cls, json_str: str) -> "BaseTypedModel":
+        """Loads a service from a JSON string.
+
+        Args:
+            json_str: the JSON string to load from.
+
+        Returns:
+            The loaded service object.
+        """
+        service_dict = json.loads(json_str)
+        class_: Type[BaseService] = source_utils.load_and_validate_class(
+            source=service_dict["type"], expected_class=BaseService
+        )
+        return class_.from_dict(service_dict)
+
     @abstractmethod
     def check_status(self) -> Tuple[ServiceState, str]:
         """Check the the current operational state of the external service.
@@ -220,7 +271,7 @@ class BaseService(BaseTypedModel, metaclass=BaseServiceMeta):
         """
 
     def update_status(self) -> None:
-        """Update the service of the service.
+        """Update the status of the service.
 
         Check the current operational state of the external service
         and update the local operational status information to reflect it.
@@ -232,21 +283,28 @@ class BaseService(BaseTypedModel, metaclass=BaseServiceMeta):
             "Running status check for service '%s' ...",
             self,
         )
-        state, err = self.check_status()
-        logger.debug(
-            "Status check results for service '%s': %s [%s]",
-            self,
-            state.name,
-            err,
-        )
-        self.status.update_state(state, err)
+        try:
+            state, err = self.check_status()
+            logger.debug(
+                "Status check results for service '%s': %s [%s]",
+                self,
+                state.name,
+                err,
+            )
+            self.status.update_state(state, err)
 
-        # don't bother checking the endpoint state if the service is not active
-        if self.status.state == ServiceState.INACTIVE:
-            return
+            # don't bother checking the endpoint state if the service is not active
+            if self.status.state == ServiceState.INACTIVE:
+                return
 
-        if self.endpoint:
-            self.endpoint.update_status()
+            if self.endpoint:
+                self.endpoint.update_status()
+        except Exception as e:
+            logger.error(
+                f"Failed to update status for service '{self}': {e}",
+                exc_info=True,
+            )
+            self.status.update_state(ServiceState.ERROR, str(e))
 
     def get_service_status_message(self) -> str:
         """Get a service status message.
@@ -373,6 +431,10 @@ class BaseService(BaseTypedModel, metaclass=BaseServiceMeta):
         """
         self.config = config
 
+    @update_service_status(
+        pre_status=ServiceState.PENDING_STARTUP,
+        post_status=ServiceState.ACTIVE,
+    )
     def start(self, timeout: int = 0) -> None:
         """Start the service and optionally wait for it to become active.
 
@@ -380,20 +442,20 @@ class BaseService(BaseTypedModel, metaclass=BaseServiceMeta):
             timeout: amount of time to wait for the service to become active.
                 If set to 0, the method will return immediately after checking
                 the service status.
-
-        Raises:
-            RuntimeError: if the service cannot be started
         """
         with console.status(f"Starting service '{self}'.\n"):
             self.admin_state = ServiceState.ACTIVE
             self.provision()
-            if timeout > 0:
-                if not self.poll_service_status(timeout):
-                    raise RuntimeError(
-                        f"Failed to start service {self}\n"
-                        + self.get_service_status_message()
-                    )
+            if timeout > 0 and not self.poll_service_status(timeout):
+                logger.error(
+                    f"Failed to start service {self}\n"
+                    + self.get_service_status_message()
+                )
 
+    @update_service_status(
+        pre_status=ServiceState.PENDING_SHUTDOWN,
+        post_status=ServiceState.INACTIVE,
+    )
     def stop(self, timeout: int = 0, force: bool = False) -> None:
         """Stop the service and optionally wait for it to shutdown.
 
@@ -403,9 +465,6 @@ class BaseService(BaseTypedModel, metaclass=BaseServiceMeta):
                 the service status.
             force: if True, the service will be stopped even if it is not
                 currently running.
-
-        Raises:
-            RuntimeError: if the service cannot be stopped
         """
         with console.status(f"Stopping service '{self}'.\n"):
             self.admin_state = ServiceState.INACTIVE
@@ -413,12 +472,40 @@ class BaseService(BaseTypedModel, metaclass=BaseServiceMeta):
             if timeout > 0:
                 self.poll_service_status(timeout)
                 if not self.is_stopped:
-                    raise RuntimeError(
+                    logger.error(
                         f"Failed to stop service {self}. Last state: "
                         f"'{self.status.state.value}'. Last error: "
                         f"'{self.status.last_error}'"
                     )
 
+    def get_prediction_url(self) -> Optional[str]:
+        """Gets the prediction URL for the endpoint.
+
+        Returns:
+            the prediction URL for the endpoint
+        """
+        prediction_url = None
+        if isinstance(self, BaseDeploymentService) and self.prediction_url:
+            prediction_url = self.prediction_url
+        elif self.endpoint:
+            prediction_url = (
+                self.endpoint.status.uri if self.endpoint.status else None
+            )
+        return prediction_url
+
+    def get_healthcheck_url(self) -> Optional[str]:
+        """Gets the healthcheck URL for the endpoint.
+
+        Returns:
+            the healthcheck URL for the endpoint
+        """
+        return (
+            self.endpoint.monitor.get_healthcheck_uri(self.endpoint)
+            if (self.endpoint and self.endpoint.monitor)
+            and isinstance(self.endpoint.monitor, HTTPEndpointHealthMonitor)
+            else None
+        )
+
     def __repr__(self) -> str:
         """String representation of the service.
 
@@ -435,14 +522,10 @@ class BaseService(BaseTypedModel, metaclass=BaseServiceMeta):
         """
         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):
@@ -456,3 +539,12 @@ class BaseDeploymentService(BaseService):
             the prediction URL for the endpoint
         """
         return None
+
+    @property
+    def healthcheck_url(self) -> Optional[str]:
+        """Gets the healthcheck URL for the endpoint.
+
+        Returns:
+            the healthcheck URL for the endpoint
+        """
+        return None

zenml/services/service_status.py

@@ -25,11 +25,12 @@ logger = get_logger(__name__)
 class ServiceState(StrEnum):
     """Possible states for the service and service endpoint."""
 
+    INACTIVE = "inactive"
     ACTIVE = "active"
     PENDING_STARTUP = "pending_startup"
-    INACTIVE = "inactive"
     PENDING_SHUTDOWN = "pending_shutdown"
     ERROR = "error"
+    SCALED_TO_ZERO = "scaled_to_zero"
 
 
 class ServiceStatus(BaseTypedModel):

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):
@@ -24,15 +24,16 @@ class ServiceType(BaseModel):
         flavor: service flavor
         name: name of the service type
         description: description of the service type
+        logo_url: logo of the service type
     """
 
     type: str
     flavor: str
     name: str = ""
     description: str = ""
+    logo_url: str = ""
 
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # make the service type immutable and hashable
-        frozen = True
+        frozen=True
+    )