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

zenml/zen_server/deploy/{local/local_provider.py → daemon/daemon_provider.py}

@@ -11,10 +11,11 @@
 #  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.
-"""Zen Server local provider implementation."""
+"""Zen Server daemon provider 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 import __version__
 from zenml.enums import ServerProviderType
@@ -31,25 +32,25 @@ 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.local.local_zen_server import (
-    LOCAL_ZENML_SERVER_DEFAULT_TIMEOUT,
+from zenml.zen_server.deploy.daemon.daemon_zen_server import (
+    DAEMON_ZENML_SERVER_DEFAULT_TIMEOUT,
     ZEN_SERVER_HEALTHCHECK_URL_PATH,
-    LocalServerDeploymentConfig,
-    LocalZenServer,
-    LocalZenServerConfig,
+    DaemonServerDeploymentConfig,
+    DaemonZenServer,
+    DaemonZenServerConfig,
 )
+from zenml.zen_server.deploy.deployment import LocalServerDeploymentConfig
 
 logger = get_logger(__name__)
 
 
-class LocalServerProvider(BaseServerProvider):
-    """Local ZenML server provider."""
+class DaemonServerProvider(BaseServerProvider):
+    """Daemon ZenML server provider."""
 
-    TYPE: ClassVar[ServerProviderType] = ServerProviderType.LOCAL
-    CONFIG_TYPE: ClassVar[
-        Type[ServerDeploymentConfig]
-    ] = LocalServerDeploymentConfig
+    TYPE: ClassVar[ServerProviderType] = ServerProviderType.DAEMON
+    CONFIG_TYPE: ClassVar[Type[LocalServerDeploymentConfig]] = (
+        DaemonServerDeploymentConfig
+    )
 
     @staticmethod
     def check_local_server_dependencies() -> None:
@@ -61,14 +62,13 @@ class LocalServerProvider(BaseServerProvider):
         try:
             # Make sure the ZenML Server dependencies are installed
             import fastapi  # noqa
-            import fastapi_utils  # noqa
             import jwt  # noqa
             import multipart  # noqa
             import uvicorn  # noqa
         except ImportError:
             # Unable to import the ZenML Server dependencies.
             raise RuntimeError(
-                "The local ZenML server provider is unavailable because the "
+                "The local daemon ZenML server provider is unavailable because the "
                 "ZenML server requirements seems to be unavailable on your machine. "
                 "This is probably because ZenML was installed without the optional "
                 "ZenML Server dependencies. To install the missing dependencies "
@@ -78,7 +78,7 @@ class LocalServerProvider(BaseServerProvider):
     @classmethod
     def _get_service_configuration(
         cls,
-        server_config: ServerDeploymentConfig,
+        server_config: LocalServerDeploymentConfig,
     ) -> Tuple[
         ServiceConfig,
         ServiceEndpointConfig,
@@ -92,13 +92,12 @@ class LocalServerProvider(BaseServerProvider):
         Returns:
             The service, service endpoint and endpoint monitor configuration.
         """
-        assert isinstance(server_config, LocalServerDeploymentConfig)
-
+        assert isinstance(server_config, DaemonServerDeploymentConfig)
         return (
-            LocalZenServerConfig(
-                root_runtime_path=LocalZenServer.config_path(),
+            DaemonZenServerConfig(
+                root_runtime_path=DaemonZenServer.config_path(),
                 singleton=True,
-                name=server_config.name,
+                name=ServerProviderType.DAEMON.value,
                 blocking=server_config.blocking,
                 server=server_config,
             ),
@@ -116,10 +115,10 @@ class LocalServerProvider(BaseServerProvider):
 
     def _create_service(
         self,
-        config: ServerDeploymentConfig,
+        config: LocalServerDeploymentConfig,
         timeout: Optional[int] = None,
     ) -> BaseService:
-        """Create, start and return the local ZenML server deployment service.
+        """Create, start and return the local daemon ZenML server deployment service.
 
         Args:
             config: The server deployment configuration.
@@ -132,16 +131,16 @@ class LocalServerProvider(BaseServerProvider):
         Raises:
             RuntimeError: If a local service is already running.
         """
-        assert isinstance(config, LocalServerDeploymentConfig)
+        assert isinstance(config, DaemonServerDeploymentConfig)
 
         if timeout is None:
-            timeout = LOCAL_ZENML_SERVER_DEFAULT_TIMEOUT
+            timeout = DAEMON_ZENML_SERVER_DEFAULT_TIMEOUT
 
         self.check_local_server_dependencies()
-        existing_service = LocalZenServer.get_service()
+        existing_service = DaemonZenServer.get_service()
         if existing_service:
             raise RuntimeError(
-                f"A local ZenML server with name '{existing_service.config.name}' "
+                f"A local daemon ZenML server with name '{existing_service.config.name}' "
                 f"is already running. Please stop it first before starting a "
                 f"new one."
             )
@@ -157,17 +156,19 @@ class LocalServerProvider(BaseServerProvider):
                 config=monitor_cfg,
             ),
         )
-        service = LocalZenServer(config=service_config, endpoint=endpoint)
+        service = DaemonZenServer(
+            uuid=uuid4(), config=service_config, endpoint=endpoint
+        )
         service.start(timeout=timeout)
         return service
 
     def _update_service(
         self,
         service: BaseService,
-        config: ServerDeploymentConfig,
+        config: LocalServerDeploymentConfig,
         timeout: Optional[int] = None,
     ) -> BaseService:
-        """Update the local ZenML server deployment service.
+        """Update the local daemon ZenML server deployment service.
 
         Args:
             service: The service instance.
@@ -179,7 +180,7 @@ class LocalServerProvider(BaseServerProvider):
             The updated service instance.
         """
         if timeout is None:
-            timeout = LOCAL_ZENML_SERVER_DEFAULT_TIMEOUT
+            timeout = DAEMON_ZENML_SERVER_DEFAULT_TIMEOUT
 
         (
             new_config,
@@ -208,7 +209,7 @@ class LocalServerProvider(BaseServerProvider):
         service: BaseService,
         timeout: Optional[int] = None,
     ) -> BaseService:
-        """Start the local ZenML server deployment service.
+        """Start the local daemon ZenML server deployment service.
 
         Args:
             service: The service instance.
@@ -219,7 +220,7 @@ class LocalServerProvider(BaseServerProvider):
             The updated service instance.
         """
         if timeout is None:
-            timeout = LOCAL_ZENML_SERVER_DEFAULT_TIMEOUT
+            timeout = DAEMON_ZENML_SERVER_DEFAULT_TIMEOUT
 
         service.start(timeout=timeout)
         return service
@@ -229,7 +230,7 @@ class LocalServerProvider(BaseServerProvider):
         service: BaseService,
         timeout: Optional[int] = None,
     ) -> BaseService:
-        """Stop the local ZenML server deployment service.
+        """Stop the local daemon ZenML server deployment service.
 
         Args:
             service: The service instance.
@@ -240,7 +241,7 @@ class LocalServerProvider(BaseServerProvider):
             The updated service instance.
         """
         if timeout is None:
-            timeout = LOCAL_ZENML_SERVER_DEFAULT_TIMEOUT
+            timeout = DAEMON_ZENML_SERVER_DEFAULT_TIMEOUT
 
         service.stop(timeout=timeout)
         return service
@@ -250,26 +251,23 @@ class LocalServerProvider(BaseServerProvider):
         service: BaseService,
         timeout: Optional[int] = None,
     ) -> None:
-        """Remove the local ZenML server deployment service.
+        """Remove the local daemon ZenML server deployment service.
 
         Args:
             service: The service instance.
             timeout: The timeout in seconds to wait until the service is
                 removed.
         """
-        assert isinstance(service, LocalZenServer)
+        assert isinstance(service, DaemonZenServer)
 
         if timeout is None:
-            timeout = LOCAL_ZENML_SERVER_DEFAULT_TIMEOUT
+            timeout = DAEMON_ZENML_SERVER_DEFAULT_TIMEOUT
 
         service.stop(timeout)
-        shutil.rmtree(LocalZenServer.config_path())
+        shutil.rmtree(DaemonZenServer.config_path())
 
-    def _get_service(self, server_name: str) -> BaseService:
-        """Get the local ZenML server deployment service.
-
-        Args:
-            server_name: The server deployment name.
+    def _get_service(self) -> BaseService:
+        """Get the local daemon ZenML server deployment service.
 
         Returns:
             The service instance.
@@ -277,31 +275,15 @@ class LocalServerProvider(BaseServerProvider):
         Raises:
             KeyError: If the server deployment is not found.
         """
-        service = LocalZenServer.get_service()
+        service = DaemonZenServer.get_service()
         if service is None:
-            raise KeyError("The local ZenML server is not deployed.")
-
-        if service.config.name != server_name:
-            raise KeyError(
-                "The local ZenML server is deployed but with a different name."
-            )
+            raise KeyError("The local daemon ZenML server is not deployed.")
 
         return service
 
-    def _list_services(self) -> List[BaseService]:
-        """Get all service instances for all deployed ZenML servers.
-
-        Returns:
-            A list of service instances.
-        """
-        service = LocalZenServer.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:
@@ -310,8 +292,8 @@ class LocalServerProvider(BaseServerProvider):
         Returns:
             The server deployment configuration.
         """
-        server = cast(LocalZenServer, service)
+        server = cast(DaemonZenServer, service)
         return server.config.server
 
 
-LocalServerProvider.register_as_provider()
+DaemonServerProvider.register_as_provider()

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

@@ -11,24 +11,31 @@
 #  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
 from typing import Dict, List, Optional, Tuple, Union, cast
 
+from pydantic import ConfigDict, Field
+
 from zenml.client import Client
 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 (
@@ -38,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.
@@ -57,30 +64,37 @@ class LocalServerDeploymentConfig(ServerDeploymentConfig):
     """
 
     port: int = 8237
-    ip_address: Union[
-        ipaddress.IPv4Address, ipaddress.IPv6Address
-    ] = ipaddress.IPv4Address(DEFAULT_LOCAL_SERVICE_IP_ADDRESS)
+    ip_address: Union[ipaddress.IPv4Address, ipaddress.IPv6Address] = Field(
+        default=ipaddress.IPv4Address(DEFAULT_LOCAL_SERVICE_IP_ADDRESS),
+        union_mode="left_to_right",
+    )
     blocking: bool = False
     store: Optional[StoreConfiguration] = None
 
-    class Config:
-        """Pydantic configuration."""
+    @property
+    def url(self) -> Optional[str]:
+        """Get the configured server URL.
 
-        extra = "forbid"
+        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
@@ -88,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
@@ -119,42 +133,19 @@ class LocalZenServer(LocalDaemonService):
         """
         return os.path.join(self.config_path(), ".zenconfig")
 
-    def _copy_global_configuration(self) -> None:
-        """Copy the global configuration to the local ZenML server location.
-
-        The local ZenML server global configuration is a copy of the local
-        global configuration. If a store configuration is explicitly set in
-        the server configuration, it will be used. Otherwise, the store
-        configuration is set to point to the local store.
-        """
-        gc = GlobalConfiguration()
-
-        # this creates a copy of the global configuration and saves it to
-        # the server configuration path. The store is set to point to the local
-        # default database unless a custom store configuration is explicitly
-        # supplied with the server configuration.
-        gc.copy_configuration(
-            config_path=self._global_config_path,
-            store_config=self.config.server.store,
-            empty_store=self.config.server.store is None,
-        )
-
     @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.
         """
-        from zenml.services import ServiceRegistry
-
         config_filename = os.path.join(cls.config_path(), "service.json")
         try:
             with open(config_filename, "r") as f:
                 return cast(
-                    LocalZenServer,
-                    ServiceRegistry().load_service_from_json(f.read()),
+                    "DaemonZenServer", DaemonZenServer.from_json(f.read())
                 )
         except FileNotFoundError:
             return None
@@ -169,22 +160,29 @@ 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 store configuration is initialized with
-        # the same path as the client.
-        env[
-            ENV_ZENML_LOCAL_STORES_PATH
-        ] = GlobalConfiguration().local_stores_path
+        # ensures that the server's default store configuration is initialized
+        # to point at the same local SQLite database as the client.
+        env[ENV_ZENML_LOCAL_STORES_PATH] = (
+            GlobalConfiguration().local_stores_path
+        )
         env[ENV_ZENML_DISABLE_DATABASE_MIGRATION] = "True"
+        env[ENV_ZENML_SERVER_AUTO_ACTIVATE] = "True"
 
         return cmd, env
 
     def provision(self) -> None:
         """Provision the service."""
-        self._copy_global_configuration()
         super().provision()
 
     def start(self, timeout: int = 0) -> None:
@@ -198,21 +196,38 @@ class LocalZenServer(LocalDaemonService):
         if not self.config.blocking:
             super().start(timeout)
         else:
-            self._copy_global_configuration()
+            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
+            # global configuration path). This is necessary to avoid polluting
+            # the client environment with the server's configuration.
             local_stores_path = GlobalConfiguration().local_stores_path
             GlobalConfiguration._reset_instance()
             Client._reset_instance()
-            config_path = os.environ.get(ENV_ZENML_CONFIG_PATH)
+            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:
-                if config_path:
-                    os.environ[ENV_ZENML_CONFIG_PATH] = config_path
+                # 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()
 
@@ -226,7 +241,7 @@ class LocalZenServer(LocalDaemonService):
         import uvicorn
 
         gc = GlobalConfiguration()
-        if gc.store and gc.store.type == StoreType.REST:
+        if gc.store_configuration.type == StoreType.REST:
             raise ValueError(
                 "The ZenML server cannot be started with REST store type."
             )
@@ -243,6 +258,7 @@ class LocalZenServer(LocalDaemonService):
                 host=self.endpoint.config.ip_address,
                 port=self.endpoint.config.port or 8000,
                 log_level="info",
+                server_header=False,
             )
         except KeyboardInterrupt:
             logger.info("ZenML Server stopped. Resuming normal execution.")