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

zenml/zen_server/deploy/deployer.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """ZenML server deployer singleton implementation."""
 
-from typing import ClassVar, Dict, Generator, List, Optional, Type, Union
+from typing import ClassVar, Dict, Generator, Optional, Type
 
 from zenml.config.global_config import GlobalConfiguration
 from zenml.enums import ServerProviderType, StoreType
@@ -21,22 +21,22 @@ from zenml.logger import get_logger
 from zenml.utils.singleton import SingletonMetaClass
 from zenml.zen_server.deploy.base_provider import BaseServerProvider
 from zenml.zen_server.deploy.deployment import (
-    ServerDeployment,
-    ServerDeploymentConfig,
+    LocalServerDeployment,
+    LocalServerDeploymentConfig,
 )
 from zenml.zen_server.deploy.exceptions import (
     ServerDeploymentError,
-    ServerDeploymentExistsError,
     ServerDeploymentNotFoundError,
     ServerProviderNotFoundError,
 )
+from zenml.zen_stores.base_zen_store import BaseZenStore
 from zenml.zen_stores.rest_zen_store import RestZenStoreConfiguration
 
 logger = get_logger(__name__)
 
 
-class ServerDeployer(metaclass=SingletonMetaClass):
-    """Server deployer singleton.
+class LocalServerDeployer(metaclass=SingletonMetaClass):
+    """Local server deployer singleton.
 
     This class is responsible for managing the various server provider
     implementations and for directing server deployment lifecycle requests to
@@ -85,318 +85,238 @@ class ServerDeployer(metaclass=SingletonMetaClass):
             )
         return cls._providers[provider_type]
 
+    def initialize_local_database(self) -> None:
+        """Initialize the local ZenML database."""
+        default_store_cfg = GlobalConfiguration().get_default_store()
+        BaseZenStore.create_store(default_store_cfg)
+
     def deploy_server(
         self,
-        config: ServerDeploymentConfig,
+        config: LocalServerDeploymentConfig,
         timeout: Optional[int] = None,
-    ) -> ServerDeployment:
-        """Deploy a new ZenML server or update an existing deployment.
+        restart: bool = False,
+    ) -> LocalServerDeployment:
+        """Deploy the local ZenML server or update the existing deployment.
 
         Args:
             config: The server deployment configuration.
             timeout: The timeout in seconds to wait until the deployment is
                 successful. If not supplied, the default timeout value specified
                 by the provider is used.
+            restart: If True, the existing server deployment will be torn down
+                and a new server will be deployed.
 
         Returns:
-            The server deployment.
+            The local server deployment.
         """
-        # We do this here to ensure that the zenml store is always initialized
-        # before the server is deployed. This is necessary because the server
-        # may require access to the local store configuration or database.
-        gc = GlobalConfiguration()
-
-        _ = gc.zen_store
+        # Ensure that the local database is always initialized before any local
+        # server is deployed or updated.
+        self.initialize_local_database()
 
         try:
-            self.get_server(config.name)
+            self.get_server()
         except ServerDeploymentNotFoundError:
             pass
         else:
-            return self.update_server(config=config, timeout=timeout)
+            return self.update_server(
+                config=config, timeout=timeout, restart=restart
+            )
 
         provider_name = config.provider.value
         provider = self.get_provider(config.provider)
 
-        logger.info(
-            f"Deploying a {provider_name} ZenML server with name "
-            f"'{config.name}'."
-        )
+        logger.info(f"Deploying a local {provider_name} ZenML server.")
         return provider.deploy_server(config, timeout=timeout)
 
     def update_server(
         self,
-        config: ServerDeploymentConfig,
+        config: LocalServerDeploymentConfig,
         timeout: Optional[int] = None,
-    ) -> ServerDeployment:
-        """Update an existing ZenML server deployment.
+        restart: bool = False,
+    ) -> LocalServerDeployment:
+        """Update an existing local ZenML server deployment.
 
         Args:
             config: The new server deployment configuration.
             timeout: The timeout in seconds to wait until the deployment is
                 successful. If not supplied, a default timeout value of 30
                 seconds is used.
+            restart: If True, the existing server deployment will be torn down
+                and a new server will be deployed.
 
         Returns:
             The updated server deployment.
-
-        Raises:
-            ServerDeploymentExistsError: If an existing deployment with the same
-                name but a different provider type is found.
         """
         # this will also raise ServerDeploymentNotFoundError if the server
         # does not exist
-        existing_server = self.get_server(config.name)
+        existing_server = self.get_server()
 
         provider = self.get_provider(config.provider)
-        existing_provider = existing_server.config.provider
-
-        if existing_provider != config.provider:
-            raise ServerDeploymentExistsError(
-                f"A server deployment with the same name '{config.name}' but "
-                f"with a different provider '{existing_provider.value}'."
-                f"is already provisioned. Please choose a different name or "
-                f"tear down the existing deployment."
+        if existing_server.config.provider != config.provider or restart:
+            existing_provider = self.get_provider(
+                existing_server.config.provider
+            )
+
+            # Tear down the existing server deployment
+            existing_provider.remove_server(
+                existing_server.config, timeout=timeout
             )
 
+            # Deploy a new server with the new provider
+            return provider.deploy_server(config, timeout=timeout)
+
         return provider.update_server(config, timeout=timeout)
 
     def remove_server(
         self,
-        server_name: str,
         timeout: Optional[int] = None,
     ) -> None:
-        """Tears down and removes all resources and files associated with a ZenML server deployment.
+        """Tears down and removes all resources and files associated with the local ZenML server deployment.
 
         Args:
-            server_name: The server deployment name.
             timeout: The timeout in seconds to wait until the deployment is
                 successfully torn down. If not supplied, a provider specific
                 default timeout value is used.
         """
         # this will also raise ServerDeploymentNotFoundError if the server
         # does not exist
-        server = self.get_server(server_name)
+        try:
+            server = self.get_server()
+        except ServerDeploymentNotFoundError:
+            return
 
         provider_name = server.config.provider.value
         provider = self.get_provider(server.config.provider)
 
-        if self.is_connected_to_server(server_name):
+        if self.is_connected_to_server():
             try:
-                self.disconnect_from_server(server_name)
+                self.disconnect_from_server()
             except Exception as e:
-                logger.warning(f"Failed to disconnect from the server: {e}")
+                logger.warning(
+                    f"Failed to disconnect from the local server: {e}"
+                )
 
-        logger.info(
-            f"Tearing down the '{server_name}' {provider_name} ZenML server."
-        )
+        logger.info(f"Tearing down the local {provider_name} ZenML server.")
         provider.remove_server(server.config, timeout=timeout)
 
-    def is_connected_to_server(self, server_name: str) -> bool:
-        """Check if the ZenML client is currently connected to a ZenML server.
-
-        Args:
-            server_name: The server deployment name.
+    def is_connected_to_server(self) -> bool:
+        """Check if the ZenML client is currently connected to the local ZenML server.
 
         Returns:
-            True if the ZenML client is connected to the ZenML server, False
+            True if the ZenML client is connected to the local ZenML server, False
             otherwise.
         """
         # this will also raise ServerDeploymentNotFoundError if the server
         # does not exist
-        server = self.get_server(server_name)
+        try:
+            server = self.get_server()
+        except ServerDeploymentNotFoundError:
+            return False
 
         gc = GlobalConfiguration()
         return (
             server.status is not None
             and server.status.url is not None
-            and gc.store is not None
-            and gc.store.url == server.status.url
+            and gc.store_configuration.url == server.status.url
         )
 
     def connect_to_server(
         self,
-        server_name: str,
-        username: str,
-        password: str,
-        verify_ssl: Union[bool, str] = True,
     ) -> None:
-        """Connect to a ZenML server instance.
-
-        Args:
-            server_name: The server deployment name.
-            username: The username to use to connect to the server.
-            password: The password to use to connect to the server.
-            verify_ssl: Either a boolean, in which case it controls whether we
-                verify the server's TLS certificate, or a string, in which case
-                it must be a path to a CA bundle to use or the CA bundle value
-                itself.
+        """Connect to the local ZenML server instance.
 
         Raises:
-            ServerDeploymentError: If the ZenML server is not running or
+            ServerDeploymentError: If the local ZenML server is not running or
                 is unreachable.
         """
         # this will also raise ServerDeploymentNotFoundError if the server
         # does not exist
-        server = self.get_server(server_name)
+        server = self.get_server()
         provider_name = server.config.provider.value
 
         gc = GlobalConfiguration()
         if not server.status or not server.status.url:
             raise ServerDeploymentError(
-                f"The {provider_name} {server_name} ZenML "
-                f"server is not currently running or is unreachable."
+                f"The local {provider_name} ZenML server is not currently "
+                "running or is unreachable."
             )
 
         store_config = RestZenStoreConfiguration(
             url=server.status.url,
-            username=username,
-            password=password,
-            verify_ssl=verify_ssl,
         )
 
-        if gc.store == store_config:
+        if gc.store_configuration == store_config:
             logger.info(
-                f"ZenML is already connected to the '{server_name}' "
+                "Your client is already connected to the local "
                 f"{provider_name} ZenML server."
             )
             return
 
         logger.info(
-            f"Connecting ZenML to the '{server_name}' "
-            f"{provider_name} ZenML server ({store_config.url})."
+            f"Connecting to the local {provider_name} ZenML server "
+            f"({store_config.url})."
         )
 
         gc.set_store(store_config)
 
         logger.info(
-            f"Connected ZenML to the '{server_name}' "
-            f"{provider_name} ZenML server ({store_config.url})."
+            f"Connected to the local {provider_name} ZenML server "
+            f"({store_config.url})."
         )
 
     def disconnect_from_server(
         self,
-        server_name: Optional[str] = None,
     ) -> None:
-        """Disconnect from a ZenML server instance.
-
-        Args:
-            server_name: The server deployment name. If supplied, the deployer
-                will check if the ZenML client is indeed connected to the server
-                and disconnect only if that is the case. Otherwise the deployer
-                will disconnect from any ZenML server.
-        """
+        """Disconnect from the ZenML server instance."""
         gc = GlobalConfiguration()
+        store_cfg = gc.store_configuration
 
-        if not gc.store or gc.store.type != StoreType.REST:
-            logger.info("ZenML is not currently connected to a ZenML server.")
-            return
-
-        if server_name:
-            # this will also raise ServerDeploymentNotFoundError if the server
-            # does not exist
-            server = self.get_server(server_name)
-            provider_name = server.config.provider.value
-
-            if not self.is_connected_to_server(server_name):
-                logger.info(
-                    f"ZenML is not currently connected to the '{server_name}' "
-                    f"{provider_name} ZenML server."
-                )
-                return
-
-            logger.info(
-                f"Disconnecting ZenML from the '{server_name}' "
-                f"{provider_name} ZenML server ({gc.store.url})."
-            )
-        else:
+        if store_cfg.type != StoreType.REST:
             logger.info(
-                f"Disconnecting ZenML from the {gc.store.url} ZenML server."
+                "Your client is not currently connected to a ZenML server."
             )
+            return
+
+        logger.info(
+            f"Disconnecting from the local ({store_cfg.url}) ZenML server."
+        )
 
         gc.set_default_store()
 
-        logger.info("Disconnected ZenML from the ZenML server.")
+        logger.info("Disconnected from the local ZenML server.")
 
     def get_server(
         self,
-        server_name: str,
-    ) -> ServerDeployment:
-        """Get a server deployment.
-
-        Args:
-            server_name: The server deployment name.
+    ) -> LocalServerDeployment:
+        """Get the local server deployment.
 
         Returns:
-            The requested server deployment.
+            The local server deployment.
 
         Raises:
-            ServerDeploymentNotFoundError: If no server deployment with the
-                given name is found.
+            ServerDeploymentNotFoundError: If no local server deployment is
+                found.
         """
         for provider in self._providers.values():
             try:
                 return provider.get_server(
-                    ServerDeploymentConfig(
-                        name=server_name, provider=provider.TYPE
-                    )
+                    LocalServerDeploymentConfig(provider=provider.TYPE)
                 )
             except ServerDeploymentNotFoundError:
                 pass
 
         raise ServerDeploymentNotFoundError(
-            f"Server deployment '{server_name}' not found."
+            "No local server deployment was found."
         )
 
-    def list_servers(
-        self,
-        server_name: Optional[str] = None,
-        provider_type: Optional[ServerProviderType] = None,
-    ) -> List[ServerDeployment]:
-        """List all server deployments.
-
-        Args:
-            server_name: The server deployment name to filter by.
-            provider_type: The server provider type to filter by.
-
-        Returns:
-            The list of server deployments.
-        """
-        providers: List[BaseServerProvider] = []
-        if provider_type:
-            providers = [self.get_provider(provider_type)]
-        else:
-            providers = list(self._providers.values())
-
-        servers: List[ServerDeployment] = []
-        for provider in providers:
-            if server_name:
-                try:
-                    servers.append(
-                        provider.get_server(
-                            ServerDeploymentConfig(
-                                name=server_name,
-                                provider=provider.TYPE,
-                            )
-                        )
-                    )
-                except ServerDeploymentNotFoundError:
-                    pass
-            else:
-                servers.extend(provider.list_servers())
-
-        return servers
-
     def get_server_logs(
         self,
-        server_name: str,
         follow: bool = False,
         tail: Optional[int] = None,
     ) -> Generator[str, bool, None]:
-        """Retrieve the logs of a ZenML server.
+        """Retrieve the logs for the local ZenML server.
 
         Args:
-            server_name: The server deployment name.
             follow: if True, the logs will be streamed as they are written
             tail: only retrieve the last NUM lines of log output.
 
@@ -405,14 +325,13 @@ class ServerDeployer(metaclass=SingletonMetaClass):
         """
         # this will also raise ServerDeploymentNotFoundError if the server
         # does not exist
-        server = self.get_server(server_name)
+        server = self.get_server()
 
         provider_name = server.config.provider.value
         provider = self.get_provider(server.config.provider)
 
         logger.info(
-            f"Fetching logs from the '{server_name}' {provider_name} ZenML "
-            f"server..."
+            f"Fetching logs from the local {provider_name} ZenML server..."
         )
         return provider.get_server_logs(
             server.config, follow=follow, tail=tail

zenml/zen_server/deploy/deployment.py

@@ -15,44 +15,50 @@
 
 from typing import Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 from zenml.enums import ServerProviderType
 from zenml.services.service_status import ServiceState
 
 
-class ServerDeploymentConfig(BaseModel):
-    """Generic server deployment configuration.
+class LocalServerDeploymentConfig(BaseModel):
+    """Generic local server deployment configuration.
 
-    All server deployment configurations should inherit from this class and
-    handle extra attributes as provider specific attributes.
+    All local server deployment configurations should inherit from this class
+    and handle extra attributes as provider specific attributes.
 
     Attributes:
-        name: Name of the server deployment.
         provider: The server provider type.
     """
 
-    name: str
     provider: ServerProviderType
 
-    class Config:
-        """Pydantic configuration class."""
+    @property
+    def url(self) -> Optional[str]:
+        """Get the configured server URL.
+
+        Returns:
+            The configured server URL.
+        """
+        return None
 
+    model_config = ConfigDict(
         # Validate attributes when assigning them. We need to set this in order
         # to have a mix of mutable and immutable attributes
-        validate_assignment = True
+        validate_assignment=True,
         # Allow extra attributes to be set in the base class. The concrete
         # classes are responsible for validating the attributes.
-        extra = "allow"
+        extra="allow",
+    )
 
 
-class ServerDeploymentStatus(BaseModel):
-    """Server deployment status.
+class LocalServerDeploymentStatus(BaseModel):
+    """Local server deployment status.
 
     Ideally this should convey the following information:
 
     * whether the server's deployment is managed by this client (i.e. if
-    the server was deployed with `zenml up`)
+    the server was deployed with `zenml login --local`)
     * for a managed deployment, the status of the deployment/tear-down, e.g.
     not deployed, deploying, running, deleting, deployment timeout/error,
     tear-down timeout/error etc.
@@ -74,7 +80,7 @@ class ServerDeploymentStatus(BaseModel):
     ca_crt: Optional[str] = None
 
 
-class ServerDeployment(BaseModel):
+class LocalServerDeployment(BaseModel):
     """Server deployment.
 
     Attributes:
@@ -82,8 +88,8 @@ class ServerDeployment(BaseModel):
         status: The server deployment status.
     """
 
-    config: ServerDeploymentConfig
-    status: Optional[ServerDeploymentStatus]
+    config: LocalServerDeploymentConfig
+    status: Optional[LocalServerDeploymentStatus] = None
 
     @property
     def is_running(self) -> bool:

zenml/zen_server/deploy/docker/docker_provider.py

@@ -14,7 +14,8 @@
 """Zen Server docker deployer implementation."""
 
 import shutil
-from typing import ClassVar, List, Optional, Tuple, Type, cast
+from typing import ClassVar, Optional, Tuple, Type, cast
+from uuid import uuid4
 
 from zenml.enums import ServerProviderType
 from zenml.logger import get_logger
@@ -30,7 +31,7 @@ from zenml.services import (
     ServiceEndpointProtocol,
 )
 from zenml.zen_server.deploy.base_provider import BaseServerProvider
-from zenml.zen_server.deploy.deployment import ServerDeploymentConfig
+from zenml.zen_server.deploy.deployment import LocalServerDeploymentConfig
 from zenml.zen_server.deploy.docker.docker_zen_server import (
     DOCKER_ZENML_SERVER_DEFAULT_TIMEOUT,
     ZEN_SERVER_HEALTHCHECK_URL_PATH,
@@ -46,14 +47,14 @@ class DockerServerProvider(BaseServerProvider):
     """Docker ZenML server provider."""
 
     TYPE: ClassVar[ServerProviderType] = ServerProviderType.DOCKER
-    CONFIG_TYPE: ClassVar[
-        Type[ServerDeploymentConfig]
-    ] = DockerServerDeploymentConfig
+    CONFIG_TYPE: ClassVar[Type[LocalServerDeploymentConfig]] = (
+        DockerServerDeploymentConfig
+    )
 
     @classmethod
     def _get_service_configuration(
         cls,
-        server_config: ServerDeploymentConfig,
+        server_config: LocalServerDeploymentConfig,
     ) -> Tuple[
         ServiceConfig,
         ServiceEndpointConfig,
@@ -74,7 +75,7 @@ class DockerServerProvider(BaseServerProvider):
                 root_runtime_path=DockerZenServer.config_path(),
                 singleton=True,
                 image=server_config.image,
-                name=server_config.name,
+                name=ServerProviderType.DOCKER.value,
                 server=server_config,
             ),
             ContainerServiceEndpointConfig(
@@ -90,7 +91,7 @@ class DockerServerProvider(BaseServerProvider):
 
     def _create_service(
         self,
-        config: ServerDeploymentConfig,
+        config: LocalServerDeploymentConfig,
         timeout: Optional[int] = None,
     ) -> BaseService:
         """Create, start and return the docker ZenML server deployment service.
@@ -131,7 +132,9 @@ class DockerServerProvider(BaseServerProvider):
                 config=monitor_cfg,
             ),
         )
-        service = DockerZenServer(config=service_config, endpoint=endpoint)
+        service = DockerZenServer(
+            uuid=uuid4(), config=service_config, endpoint=endpoint
+        )
 
         service.start(timeout=timeout)
         return service
@@ -139,7 +142,7 @@ class DockerServerProvider(BaseServerProvider):
     def _update_service(
         self,
         service: BaseService,
-        config: ServerDeploymentConfig,
+        config: LocalServerDeploymentConfig,
         timeout: Optional[int] = None,
     ) -> BaseService:
         """Update the docker ZenML server deployment service.
@@ -241,12 +244,9 @@ class DockerServerProvider(BaseServerProvider):
         service.stop(timeout)
         shutil.rmtree(DockerZenServer.config_path())
 
-    def _get_service(self, server_name: str) -> BaseService:
+    def _get_service(self) -> BaseService:
         """Get the docker ZenML server deployment service.
 
-        Args:
-            server_name: The server deployment name.
-
         Returns:
             The service instance.
 
@@ -257,27 +257,11 @@ class DockerServerProvider(BaseServerProvider):
         if service is None:
             raise KeyError("The docker ZenML server is not deployed.")
 
-        if service.config.name != server_name:
-            raise KeyError(
-                "The docker ZenML server is deployed but with a different name."
-            )
-
         return service
 
-    def _list_services(self) -> List[BaseService]:
-        """Get all service instances for all deployed ZenML servers.
-
-        Returns:
-            A list of service instances.
-        """
-        service = DockerZenServer.get_service()
-        if service:
-            return [service]
-        return []
-
     def _get_deployment_config(
         self, service: BaseService
-    ) -> ServerDeploymentConfig:
+    ) -> LocalServerDeploymentConfig:
         """Recreate the server deployment configuration from a service instance.
 
         Args: