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

zenml/models/v2/misc/server_models.py

@@ -13,6 +13,8 @@
 #  permissions and limitations under the License.
 """Model definitions for ZenML servers."""
 
+from datetime import datetime
+from typing import Dict, Optional
 from uuid import UUID, uuid4
 
 from pydantic import BaseModel, Field
@@ -50,10 +52,16 @@ class ServerModel(BaseModel):
 
     id: UUID = Field(default_factory=uuid4, title="The unique server id.")
 
+    name: Optional[str] = Field(None, title="The name of the ZenML server.")
+
     version: str = Field(
         title="The ZenML version that the server is running.",
     )
 
+    active: bool = Field(
+        True, title="Flag to indicate whether the server is active."
+    )
+
     debug: bool = Field(
         False, title="Flag to indicate whether ZenML is running on debug mode."
     )
@@ -73,6 +81,67 @@ class ServerModel(BaseModel):
     auth_scheme: AuthScheme = Field(
         title="The authentication scheme that the server is using.",
     )
+    server_url: str = Field(
+        "",
+        title="The URL where the ZenML server API is reachable. If not "
+        "specified, the clients will use the same URL used to connect them to "
+        "the ZenML server.",
+    )
+    dashboard_url: str = Field(
+        "",
+        title="The URL where the ZenML dashboard is reachable. If "
+        "not specified, the `server_url` value will be used instead.",
+    )
+    analytics_enabled: bool = Field(
+        default=True,  # We set a default for migrations from < 0.57.0
+        title="Enable server-side analytics.",
+    )
+
+    metadata: Dict[str, str] = Field(
+        {},
+        title="The metadata associated with the server.",
+    )
+
+    last_user_activity: Optional[datetime] = Field(
+        None,
+        title="Timestamp of latest user activity traced on the server.",
+    )
+
+    pro_dashboard_url: Optional[str] = Field(
+        None,
+        title="The base URL of the ZenML Pro dashboard to which the server "
+        "is connected. Only set if the server is a ZenML Pro server.",
+    )
+
+    pro_api_url: Optional[str] = Field(
+        None,
+        title="The base URL of the ZenML Pro API to which the server is "
+        "connected. Only set if the server is a ZenML Pro server.",
+    )
+
+    pro_organization_id: Optional[UUID] = Field(
+        None,
+        title="The ID of the ZenML Pro organization to which the server is "
+        "connected. Only set if the server is a ZenML Pro server.",
+    )
+
+    pro_organization_name: Optional[str] = Field(
+        None,
+        title="The name of the ZenML Pro organization to which the server is "
+        "connected. Only set if the server is a ZenML Pro server.",
+    )
+
+    pro_tenant_id: Optional[UUID] = Field(
+        None,
+        title="The ID of the ZenML Pro tenant to which the server is connected. "
+        "Only set if the server is a ZenML Pro server.",
+    )
+
+    pro_tenant_name: Optional[str] = Field(
+        None,
+        title="The name of the ZenML Pro tenant to which the server is connected. "
+        "Only set if the server is a ZenML Pro server.",
+    )
 
     def is_local(self) -> bool:
         """Return whether the server is running locally.
@@ -85,3 +154,34 @@ class ServerModel(BaseModel):
         # Local ZenML servers are identifiable by the fact that their
         # server ID is the same as the local client (user) ID.
         return self.id == GlobalConfiguration().user_id
+
+    def is_pro_server(self) -> bool:
+        """Return whether the server is a ZenML Pro server.
+
+        Returns:
+            True if the server is a ZenML Pro server, False otherwise.
+        """
+        return self.deployment_type == ServerDeploymentType.CLOUD
+
+
+class ServerLoadInfo(BaseModel):
+    """Domain model for ZenML server load information."""
+
+    threads: int = Field(
+        title="Number of threads that the server is currently using."
+    )
+
+    db_connections_total: int = Field(
+        title="Total number of database connections (active and idle) that the "
+        "server currently has established."
+    )
+
+    db_connections_active: int = Field(
+        title="Number of database connections that the server is currently "
+        "actively using to make queries or transactions."
+    )
+
+    db_connections_overflow: int = Field(
+        title="Number of overflow database connections that the server is "
+        "currently actively using to make queries or transactions."
+    )

zenml/models/v2/misc/service_connector_type.py

@@ -12,11 +12,11 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Model definitions for ZenML service connectors."""
-import json
+
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
 from uuid import UUID
 
-from pydantic import BaseModel, Field, validator
+from pydantic import BaseModel, Field, field_validator
 
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.logger import get_logger
@@ -140,7 +140,8 @@ class AuthenticationMethodModel(BaseModel):
             **values: The data to initialize the authentication method with.
         """
         if config_class:
-            values["config_schema"] = json.loads(config_class.schema_json())
+            values["config_schema"] = config_class.model_json_schema()
+
         super().__init__(**values)
         self._config_class = config_class
 
@@ -214,11 +215,6 @@ class AuthenticationMethodModel(BaseModel):
 
         return expiration_seconds
 
-    class Config:
-        """Pydantic config class."""
-
-        underscore_attrs_are_private = True
-
 
 class ServiceConnectorTypeModel(BaseModel):
     """Service connector type specification.
@@ -332,7 +328,8 @@ class ServiceConnectorTypeModel(BaseModel):
         """
         self._connector_class = connector_class
 
-    @validator("resource_types")
+    @field_validator("resource_types")
+    @classmethod
     def validate_resource_types(
         cls, values: List[ResourceTypeModel]
     ) -> List[ResourceTypeModel]:
@@ -359,7 +356,8 @@ class ServiceConnectorTypeModel(BaseModel):
 
         return values
 
-    @validator("auth_methods")
+    @field_validator("auth_methods")
+    @classmethod
     def validate_auth_methods(
         cls, values: List[AuthenticationMethodModel]
     ) -> List[AuthenticationMethodModel]:
@@ -471,11 +469,6 @@ class ServiceConnectorTypeModel(BaseModel):
 
         return auth_method_spec, resource_type_spec
 
-    class Config:
-        """Pydantic config class."""
-
-        underscore_attrs_are_private = True
-
 
 class ServiceConnectorRequirements(BaseModel):
     """Service connector requirements.
@@ -599,8 +592,7 @@ class ServiceConnectorResourcesModel(BaseModel):
     )
 
     connector_type: Union[str, "ServiceConnectorTypeModel"] = Field(
-        title="The type of service connector.",
-        max_length=STR_FIELD_MAX_LENGTH,
+        title="The type of service connector.", union_mode="left_to_right"
     )
 
     resources: List[ServiceConnectorTypedResourcesModel] = Field(

zenml/models/v2/misc/stack_deployment.py

@@ -0,0 +1,96 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  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.
+"""Models related to cloud stack deployments."""
+
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel, Field
+
+from zenml.enums import StackDeploymentProvider
+from zenml.models.v2.core.service_connector import ServiceConnectorResponse
+from zenml.models.v2.core.stack import StackResponse
+
+
+class StackDeploymentInfo(BaseModel):
+    """Information about a stack deployment."""
+
+    provider: StackDeploymentProvider = Field(
+        title="The provider of the stack deployment."
+    )
+    description: str = Field(
+        title="The description of the stack deployment.",
+        description="The description of the stack deployment.",
+    )
+    instructions: str = Field(
+        title="The instructions for deploying the stack.",
+        description="The instructions for deploying the stack.",
+    )
+    post_deploy_instructions: str = Field(
+        title="The instructions for post-deployment.",
+        description="The instructions for post-deployment.",
+    )
+    integrations: List[str] = Field(
+        title="ZenML integrations required for the stack.",
+        description="The list of ZenML integrations that need to be installed "
+        "for the stack to be usable.",
+    )
+    permissions: Dict[str, List[str]] = Field(
+        title="The permissions granted to ZenML to access the cloud resources.",
+        description="The permissions granted to ZenML to access the cloud "
+        "resources, as a dictionary grouping permissions by resource.",
+    )
+    locations: Dict[str, str] = Field(
+        title="The locations where the stack can be deployed.",
+        description="The locations where the stack can be deployed, as a "
+        "dictionary mapping location names to descriptions.",
+    )
+    skypilot_default_regions: Dict[str, str] = Field(
+        title="The locations where the Skypilot clusters can be deployed by default.",
+        description="The locations where the Skypilot clusters can be deployed by default, as a "
+        "dictionary mapping location names to descriptions.",
+    )
+
+
+class StackDeploymentConfig(BaseModel):
+    """Configuration about a stack deployment."""
+
+    deployment_url: str = Field(
+        title="The cloud provider console URL where the stack will be deployed.",
+    )
+    deployment_url_text: str = Field(
+        title="A textual description for the cloud provider console URL.",
+    )
+    configuration: Optional[str] = Field(
+        default=None,
+        title="Configuration for the stack deployment that the user must "
+        "manually configure into the cloud provider console.",
+    )
+    instructions: Optional[str] = Field(
+        default=None,
+        title="Instructions for deploying the stack.",
+    )
+
+
+class DeployedStack(BaseModel):
+    """Information about a deployed stack."""
+
+    stack: StackResponse = Field(
+        title="The stack that was deployed.",
+        description="The stack that was deployed.",
+    )
+    service_connector: Optional[ServiceConnectorResponse] = Field(
+        default=None,
+        title="The service connector for the deployed stack.",
+        description="The service connector for the deployed stack.",
+    )

zenml/models/v2/misc/user_auth.py

@@ -22,6 +22,7 @@ from pydantic import Field, SecretStr
 
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.models.v2.base.base import BaseZenModel
+from zenml.utils.secret_utils import PlainSerializedSecretStr
 
 if TYPE_CHECKING:
     from passlib.context import CryptContext
@@ -47,8 +48,12 @@ class UserAuthModel(BaseZenModel):
         "account."
     )
 
-    activation_token: Optional[SecretStr] = Field(default=None, exclude=True)
-    password: Optional[SecretStr] = Field(default=None, exclude=True)
+    activation_token: Optional[PlainSerializedSecretStr] = Field(
+        default=None, exclude=True
+    )
+    password: Optional[PlainSerializedSecretStr] = Field(
+        default=None, exclude=True
+    )
     name: str = Field(
         title="The unique username for the account.",
         max_length=STR_FIELD_MAX_LENGTH,
@@ -68,13 +73,6 @@ class UserAuthModel(BaseZenModel):
         "`false` if skipped.",
     )
 
-    hub_token: Optional[str] = Field(
-        default=None,
-        title="JWT Token for the connected Hub account. Only relevant for user "
-        "accounts.",
-        max_length=STR_FIELD_MAX_LENGTH,
-    )
-
     @classmethod
     def _get_crypt_context(cls) -> "CryptContext":
         """Returns the password encryption context.

zenml/orchestrators/__init__.py

@@ -30,6 +30,9 @@ from zenml.orchestrators.base_orchestrator import (
 from zenml.orchestrators.containerized_orchestrator import (
     ContainerizedOrchestrator,
 )
+from zenml.orchestrators.wheeled_orchestrator import (
+    WheeledOrchestrator,
+)
 from zenml.orchestrators.local.local_orchestrator import (
     LocalOrchestrator,
     LocalOrchestratorFlavor,
@@ -44,6 +47,7 @@ __all__ = [
     "BaseOrchestratorConfig",
     "BaseOrchestratorFlavor",
     "ContainerizedOrchestrator",
+    "WheeledOrchestrator",
     "LocalOrchestrator",
     "LocalOrchestratorFlavor",
     "LocalDockerOrchestrator",

zenml/orchestrators/base_orchestrator.py

@@ -12,20 +12,29 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Base orchestrator class."""
+
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Dict, Optional, Type, cast
+from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional, Type, cast
+from uuid import UUID
 
-from pydantic import root_validator
+from pydantic import model_validator
 
-from zenml.enums import StackComponentType
+from zenml.constants import (
+    ENV_ZENML_PREVENT_CLIENT_SIDE_CACHING,
+    handle_bool_env_var,
+)
+from zenml.enums import ExecutionStatus, StackComponentType
 from zenml.logger import get_logger
+from zenml.metadata.metadata_types import MetadataType
+from zenml.orchestrators.publish_utils import publish_pipeline_run_metadata
 from zenml.orchestrators.step_launcher import StepLauncher
 from zenml.orchestrators.utils import get_config_environment_vars
 from zenml.stack import Flavor, Stack, StackComponent, StackComponentConfig
+from zenml.utils.pydantic_utils import before_validator_handler
 
 if TYPE_CHECKING:
     from zenml.config.step_configurations import Step
-    from zenml.models import PipelineDeploymentResponse
+    from zenml.models import PipelineDeploymentResponse, PipelineRunResponse
 
 logger = get_logger(__name__)
 
@@ -33,27 +42,56 @@ logger = get_logger(__name__)
 class BaseOrchestratorConfig(StackComponentConfig):
     """Base orchestrator config."""
 
-    @root_validator(pre=True)
-    def _deprecations(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def _deprecations(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         """Validate and/or remove deprecated fields.
 
         Args:
-            values: The values to validate.
+            data: The values to validate.
 
         Returns:
             The validated values.
         """
-        if "custom_docker_base_image_name" in values:
-            image_name = values.pop("custom_docker_base_image_name", None)
+        if "custom_docker_base_image_name" in data:
+            image_name = data.pop("custom_docker_base_image_name", None)
             if image_name:
                 logger.warning(
                     "The 'custom_docker_base_image_name' field has been "
                     "deprecated. To use a custom base container image with your "
                     "orchestrators, please use the DockerSettings in your "
-                    "pipeline (see https://docs.zenml.io/user-guide/advanced-guide/environment-management/containerize-your-pipeline)."
+                    "pipeline (see https://docs.zenml.io/how-to/infrastructure-deployment/customize-docker-builds)."
                 )
 
-        return values
+        return data
+
+    @property
+    def is_synchronous(self) -> bool:
+        """Whether the orchestrator runs synchronous or not.
+
+        Returns:
+            Whether the orchestrator runs synchronous or not.
+        """
+        return False
+
+    @property
+    def is_schedulable(self) -> bool:
+        """Whether the orchestrator is schedulable or not.
+
+        Returns:
+            Whether the orchestrator is schedulable or not.
+        """
+        return False
+
+    @property
+    def supports_client_side_caching(self) -> bool:
+        """Whether the orchestrator supports client side caching.
+
+        Returns:
+            Whether the orchestrator supports client side caching.
+        """
+        return True
 
 
 class BaseOrchestrator(StackComponent, ABC):
@@ -102,7 +140,7 @@ class BaseOrchestrator(StackComponent, ABC):
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
-    ) -> Any:
+    ) -> Optional[Iterator[Dict[str, MetadataType]]]:
         """The method needs to be implemented by the respective orchestrator.
 
         Depending on the type of orchestrator you'll have to perform slightly
@@ -138,38 +176,96 @@ class BaseOrchestrator(StackComponent, ABC):
             environment: Environment variables to set in the orchestration
                 environment. These don't need to be set if running locally.
 
-        Returns:
-            The optional return value from this method will be returned by the
-            `pipeline_instance.run()` call when someone is running a pipeline.
+        Yields:
+            Metadata for the pipeline run.
         """
 
     def run(
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
-    ) -> Any:
+        placeholder_run: Optional["PipelineRunResponse"] = None,
+    ) -> None:
         """Runs a pipeline on a stack.
 
         Args:
             deployment: The pipeline deployment.
             stack: The stack on which to run the pipeline.
-
-        Returns:
-            Orchestrator-specific return value.
+            placeholder_run: An optional placeholder run for the deployment.
+                This will be deleted in case the pipeline deployment failed.
         """
         self._prepare_run(deployment=deployment)
 
-        environment = get_config_environment_vars(deployment=deployment)
+        pipeline_run_id: Optional[UUID] = None
+        schedule_id: Optional[UUID] = None
+        if deployment.schedule:
+            schedule_id = deployment.schedule.id
+        if placeholder_run:
+            pipeline_run_id = placeholder_run.id
 
-        try:
-            result = self.prepare_or_run_pipeline(
-                deployment=deployment, stack=stack, environment=environment
+        environment = get_config_environment_vars(
+            schedule_id=schedule_id,
+            pipeline_run_id=pipeline_run_id,
+        )
+
+        prevent_client_side_caching = handle_bool_env_var(
+            ENV_ZENML_PREVENT_CLIENT_SIDE_CACHING, default=False
+        )
+
+        if (
+            placeholder_run
+            and self.config.supports_client_side_caching
+            and not deployment.schedule
+            and not prevent_client_side_caching
+        ):
+            from zenml.orchestrators import step_run_utils
+
+            cached_invocations = step_run_utils.create_cached_step_runs(
+                deployment=deployment,
+                pipeline_run=placeholder_run,
+                stack=stack,
             )
+
+            for invocation_id in cached_invocations:
+                # Remove the cached step invocations from the deployment so
+                # the orchestrator does not try to run them
+                deployment.step_configurations.pop(invocation_id)
+
+            for step in deployment.step_configurations.values():
+                for invocation_id in cached_invocations:
+                    if invocation_id in step.spec.upstream_steps:
+                        step.spec.upstream_steps.remove(invocation_id)
+
+            if len(deployment.step_configurations) == 0:
+                # All steps were cached, we update the pipeline run status and
+                # don't actually use the orchestrator to run the pipeline
+                self._cleanup_run()
+                logger.info("All steps of the pipeline run were cached.")
+                return
+        else:
+            logger.debug("Skipping client-side caching.")
+
+        try:
+            if metadata_iterator := self.prepare_or_run_pipeline(
+                deployment=deployment,
+                stack=stack,
+                environment=environment,
+            ):
+                for metadata_dict in metadata_iterator:
+                    try:
+                        if placeholder_run:
+                            publish_pipeline_run_metadata(
+                                pipeline_run_id=placeholder_run.id,
+                                pipeline_run_metadata={self.id: metadata_dict},
+                            )
+                    except Exception as e:
+                        logger.debug(
+                            "Something went went wrong trying to publish the"
+                            f"run metadata: {e}"
+                        )
         finally:
             self._cleanup_run()
 
-        return result
-
     def run_step(self, step: "Step") -> None:
         """Runs the given step.
 
@@ -217,6 +313,21 @@ class BaseOrchestrator(StackComponent, ABC):
         """Cleans up the active run."""
         self._active_deployment = None
 
+    def fetch_status(self, run: "PipelineRunResponse") -> ExecutionStatus:
+        """Refreshes the status of a specific pipeline run.
+
+        Args:
+            run: A pipeline run response to fetch its status.
+
+        Raises:
+            NotImplementedError: If any orchestrator inheriting from the base
+                class does not implement this logic.
+        """
+        raise NotImplementedError(
+            "The fetch status functionality is not implemented for the "
+            f"'{self.__class__.__name__}' orchestrator."
+        )
+
 
 class BaseOrchestratorFlavor(Flavor):
     """Base orchestrator flavor class."""

zenml/orchestrators/containerized_orchestrator.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Containerized orchestrator class."""
+
 from abc import ABC
 from typing import List, Optional
 

zenml/orchestrators/dag_runner.py

@@ -14,6 +14,7 @@
 """DAG (Directed Acyclic Graph) Runners."""
 
 import threading
+import time
 from collections import defaultdict
 from enum import Enum
 from typing import Any, Callable, Dict, List
@@ -66,7 +67,10 @@ class ThreadedDagRunner:
     """
 
     def __init__(
-        self, dag: Dict[str, List[str]], run_fn: Callable[[str], Any]
+        self,
+        dag: Dict[str, List[str]],
+        run_fn: Callable[[str], Any],
+        parallel_node_startup_waiting_period: float = 0.0,
     ) -> None:
         """Define attributes and initialize all nodes in waiting state.
 
@@ -75,7 +79,12 @@ class ThreadedDagRunner:
                 E.g.: [(1->2), (1->3), (2->4), (3->4)] should be represented as
                 `dag={2: [1], 3: [1], 4: [2, 3]}`
             run_fn: A function `run_fn(node)` that runs a single node
+            parallel_node_startup_waiting_period: Delay in seconds to wait in
+                between starting parallel nodes.
         """
+        self.parallel_node_startup_waiting_period = (
+            parallel_node_startup_waiting_period
+        )
         self.dag = dag
         self.reversed_dag = reverse_dag(dag)
         self.run_fn = run_fn
@@ -154,9 +163,12 @@ class ThreadedDagRunner:
             self.node_states[node] = NodeStatus.COMPLETED
 
         # Run downstream nodes.
-        threads = []
+        threads: List[threading.Thread] = []
         for downstram_node in self.reversed_dag[node]:
             if self._can_run(downstram_node):
+                if threads and self.parallel_node_startup_waiting_period > 0:
+                    time.sleep(self.parallel_node_startup_waiting_period)
+
                 thread = self._run_node_in_thread(downstram_node)
                 threads.append(thread)
 
@@ -173,9 +185,12 @@ class ThreadedDagRunner:
         # Run all nodes that can be started immediately.
         # These will, in turn, start other nodes once all of their respective
         # upstream nodes have completed.
-        threads = []
+        threads: List[threading.Thread] = []
         for node in self.nodes:
             if self._can_run(node):
+                if threads and self.parallel_node_startup_waiting_period > 0:
+                    time.sleep(self.parallel_node_startup_waiting_period)
+
                 thread = self._run_node_in_thread(node)
                 threads.append(thread)