zenml-nightly 0.68.0.dev20241027__py3-none-any.whl → 0.68.1.dev20241101__py3-none-any.whl

zenml/zen_server/deploy/{local/local_zen_server.py → daemon/daemon_zen_server.py}

@@ -11,7 +11,7 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
-"""Local ZenML server deployment service implementation."""
+"""Local daemon ZenML server deployment service implementation."""
 
 import ipaddress
 import os
@@ -24,14 +24,18 @@ from zenml.config.global_config import GlobalConfiguration
 from zenml.config.store_config import StoreConfiguration
 from zenml.constants import (
     DEFAULT_LOCAL_SERVICE_IP_ADDRESS,
+    ENV_ZENML_ANALYTICS_OPT_IN,
     ENV_ZENML_CONFIG_PATH,
     ENV_ZENML_DISABLE_DATABASE_MIGRATION,
     ENV_ZENML_LOCAL_STORES_PATH,
+    ENV_ZENML_SERVER,
+    ENV_ZENML_SERVER_AUTH_SCHEME,
     ENV_ZENML_SERVER_AUTO_ACTIVATE,
     ENV_ZENML_SERVER_DEPLOYMENT_TYPE,
+    ENV_ZENML_USER_ID,
     ZEN_SERVER_ENTRYPOINT,
 )
-from zenml.enums import StoreType
+from zenml.enums import AuthScheme, StoreType
 from zenml.logger import get_logger
 from zenml.models import ServerDeploymentType
 from zenml.services import (
@@ -41,16 +45,16 @@ from zenml.services import (
     ServiceType,
 )
 from zenml.utils.io_utils import get_global_config_directory
-from zenml.zen_server.deploy.deployment import ServerDeploymentConfig
+from zenml.zen_server.deploy.deployment import LocalServerDeploymentConfig
 
 logger = get_logger(__name__)
 
 ZEN_SERVER_HEALTHCHECK_URL_PATH = "health"
-LOCAL_ZENML_SERVER_DEFAULT_TIMEOUT = 30
+DAEMON_ZENML_SERVER_DEFAULT_TIMEOUT = 30
 
 
-class LocalServerDeploymentConfig(ServerDeploymentConfig):
-    """Local server deployment configuration.
+class DaemonServerDeploymentConfig(LocalServerDeploymentConfig):
+    """Daemon server deployment configuration.
 
     Attributes:
         port: The TCP port number where the server is accepting connections.
@@ -67,21 +71,30 @@ class LocalServerDeploymentConfig(ServerDeploymentConfig):
     blocking: bool = False
     store: Optional[StoreConfiguration] = None
 
+    @property
+    def url(self) -> Optional[str]:
+        """Get the configured server URL.
+
+        Returns:
+            The configured server URL.
+        """
+        return f"http://{self.ip_address}:{self.port}"
+
     model_config = ConfigDict(extra="forbid")
 
 
-class LocalZenServerConfig(LocalDaemonServiceConfig):
-    """Local Zen server configuration.
+class DaemonZenServerConfig(LocalDaemonServiceConfig):
+    """Local daemon Zen server configuration.
 
     Attributes:
         server: The deployment configuration.
     """
 
-    server: LocalServerDeploymentConfig
+    server: DaemonServerDeploymentConfig
 
 
-class LocalZenServer(LocalDaemonService):
-    """Service daemon that can be used to start a local ZenML Server.
+class DaemonZenServer(LocalDaemonService):
+    """Service daemon that can be used to start a local daemon ZenML server.
 
     Attributes:
         config: service configuration
@@ -89,26 +102,26 @@ class LocalZenServer(LocalDaemonService):
     """
 
     SERVICE_TYPE = ServiceType(
-        name="local_zenml_server",
+        name="daemon_zenml_server",
         type="zen_server",
-        flavor="local",
-        description="Local ZenML server deployment",
+        flavor="daemon",
+        description="local daemon ZenML server deployment",
     )
 
-    config: LocalZenServerConfig
+    config: DaemonZenServerConfig
     endpoint: LocalDaemonServiceEndpoint
 
     @classmethod
     def config_path(cls) -> str:
-        """Path to the directory where the local ZenML server files are located.
+        """Path to the directory where the local daemon ZenML server files are located.
 
         Returns:
-            Path to the local ZenML server runtime directory.
+            Path to the local daemon ZenML server runtime directory.
         """
         return os.path.join(
             get_global_config_directory(),
             "zen_server",
-            "local",
+            "daemon",
         )
 
     @property
@@ -121,18 +134,18 @@ class LocalZenServer(LocalDaemonService):
         return os.path.join(self.config_path(), ".zenconfig")
 
     @classmethod
-    def get_service(cls) -> Optional["LocalZenServer"]:
-        """Load and return the local ZenML server service, if present.
+    def get_service(cls) -> Optional["DaemonZenServer"]:
+        """Load and return the local daemon ZenML server service, if present.
 
         Returns:
-            The local ZenML server service or None, if the local server
+            The local daemon ZenML server service or None, if the local server
             deployment is not found.
         """
         config_filename = os.path.join(cls.config_path(), "service.json")
         try:
             with open(config_filename, "r") as f:
                 return cast(
-                    "LocalZenServer", LocalZenServer.from_json(f.read())
+                    "DaemonZenServer", DaemonZenServer.from_json(f.read())
                 )
         except FileNotFoundError:
             return None
@@ -147,8 +160,15 @@ class LocalZenServer(LocalDaemonService):
             The command to start the daemon and the environment variables to
             set for the command.
         """
+        gc = GlobalConfiguration()
+
         cmd, env = super()._get_daemon_cmd()
+        env[ENV_ZENML_SERVER] = "true"
         env[ENV_ZENML_CONFIG_PATH] = self._global_config_path
+        env[ENV_ZENML_ANALYTICS_OPT_IN] = str(gc.analytics_opt_in)
+        env[ENV_ZENML_USER_ID] = str(gc.user_id)
+        # Disable authentication for the local server
+        env[ENV_ZENML_SERVER_AUTH_SCHEME] = AuthScheme.NO_AUTH.value
         env[ENV_ZENML_SERVER_DEPLOYMENT_TYPE] = ServerDeploymentType.LOCAL
         # Set the local stores path to the same path used by the client. This
         # ensures that the server's default store configuration is initialized
@@ -176,6 +196,8 @@ class LocalZenServer(LocalDaemonService):
         if not self.config.blocking:
             super().start(timeout)
         else:
+            gc = GlobalConfiguration()
+
             # In the blocking mode, we need to temporarily set the environment
             # variables for the running process to make it look like the server
             # is running in a separate environment (i.e. using a different
@@ -185,21 +207,27 @@ class LocalZenServer(LocalDaemonService):
             GlobalConfiguration._reset_instance()
             Client._reset_instance()
             original_config_path = os.environ.get(ENV_ZENML_CONFIG_PATH)
+            os.environ[ENV_ZENML_SERVER] = "true"
             os.environ[ENV_ZENML_CONFIG_PATH] = self._global_config_path
+            os.environ[ENV_ZENML_ANALYTICS_OPT_IN] = str(gc.analytics_opt_in)
+            os.environ[ENV_ZENML_USER_ID] = str(gc.user_id)
             # Set the local stores path to the same path used by the client.
             # This ensures that the server's default store configuration is
             # initialized to point at the same local SQLite database as the
             # client.
             os.environ[ENV_ZENML_LOCAL_STORES_PATH] = local_stores_path
+            os.environ[ENV_ZENML_SERVER_AUTH_SCHEME] = AuthScheme.NO_AUTH.value
             try:
                 self.run()
             finally:
                 # Restore the original client environment variables
+                del os.environ[ENV_ZENML_SERVER]
                 if original_config_path:
                     os.environ[ENV_ZENML_CONFIG_PATH] = original_config_path
                 else:
                     del os.environ[ENV_ZENML_CONFIG_PATH]
                 del os.environ[ENV_ZENML_LOCAL_STORES_PATH]
+                del os.environ[ENV_ZENML_SERVER_AUTH_SCHEME]
                 GlobalConfiguration._reset_instance()
                 Client._reset_instance()
 

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,126 +85,133 @@ 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 (
@@ -215,188 +222,101 @@ class ServerDeployer(metaclass=SingletonMetaClass):
 
     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_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 store_cfg.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 ({store_cfg.url})."
-            )
-        else:
             logger.info(
-                f"Disconnecting ZenML from the {store_cfg.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 " f"server..."
         )
         return provider.get_server_logs(
             server.config, follow=follow, tail=tail