zenml-nightly 0.58.2.dev20240623__py3-none-any.whl → 0.61.0.dev20240712__py3-none-any.whl

zenml/logging/step_logging.py

@@ -145,61 +145,64 @@ def fetch_logs(
         )
 
     artifact_store = _load_artifact_store(artifact_store_id, zen_store)
-    if not artifact_store.isdir(logs_uri):
-        return _read_file(logs_uri, offset, length)
-    else:
-        files = artifact_store.listdir(logs_uri)
-        if len(files) == 1:
-            return _read_file(
-                os.path.join(logs_uri, str(files[0])), offset, length
-            )
+    try:
+        if not artifact_store.isdir(logs_uri):
+            return _read_file(logs_uri, offset, length)
         else:
-            is_negative_offset = offset < 0
-            files.sort(reverse=is_negative_offset)
-
-            # search for the first file we need to read
-            latest_file_id = 0
-            for i, file in enumerate(files):
-                file_size: int = artifact_store.size(
-                    os.path.join(logs_uri, str(file))
-                )  # type: ignore[assignment]
-
-                if is_negative_offset:
-                    if file_size >= -offset:
-                        latest_file_id = -(i + 1)
-                        break
+            files = artifact_store.listdir(logs_uri)
+            if len(files) == 1:
+                return _read_file(
+                    os.path.join(logs_uri, str(files[0])), offset, length
+                )
+            else:
+                is_negative_offset = offset < 0
+                files.sort(reverse=is_negative_offset)
+
+                # search for the first file we need to read
+                latest_file_id = 0
+                for i, file in enumerate(files):
+                    file_size: int = artifact_store.size(
+                        os.path.join(logs_uri, str(file))
+                    )  # type: ignore[assignment]
+
+                    if is_negative_offset:
+                        if file_size >= -offset:
+                            latest_file_id = -(i + 1)
+                            break
+                        else:
+                            offset += file_size
                     else:
-                        offset += file_size
-                else:
-                    if file_size > offset:
-                        latest_file_id = i
+                        if file_size > offset:
+                            latest_file_id = i
+                            break
+                        else:
+                            offset -= file_size
+
+                # read the files according to pre-filtering
+                files.sort()
+                ret = []
+                for file in files[latest_file_id:]:
+                    ret.append(
+                        _read_file(
+                            os.path.join(logs_uri, str(file)),
+                            offset,
+                            length,
+                        )
+                    )
+                    offset = 0
+                    length -= len(ret[-1])
+                    if length <= 0:
+                        # stop further reading, if the whole length is already read
                         break
-                    else:
-                        offset -= file_size
-
-            # read the files according to pre-filtering
-            files.sort()
-            ret = []
-            for file in files[latest_file_id:]:
-                ret.append(
-                    _read_file(
-                        os.path.join(logs_uri, str(file)),
-                        offset,
-                        length,
+
+                if not ret:
+                    raise DoesNotExistException(
+                        f"Folder '{logs_uri}' is empty in artifact store "
+                        f"'{artifact_store.name}'."
                     )
-                )
-                offset = 0
-                length -= len(ret[-1])
-                if length <= 0:
-                    # stop further reading, if the whole length is already read
-                    break
-
-            if not ret:
-                raise DoesNotExistException(
-                    f"Folder '{logs_uri}' is empty in artifact store "
-                    f"'{artifact_store.name}'."
-                )
-            return "".join(ret)
+                return "".join(ret)
+    finally:
+        artifact_store.cleanup()
 
 
 class StepLogsStorage:

zenml/models/__init__.py

@@ -51,6 +51,15 @@ from zenml.models.v2.base.filter import (
 from zenml.models.v2.base.page import Page
 
 # V2 Core
+from zenml.models.v2.core.action import (
+    ActionFilter,
+    ActionRequest,
+    ActionResponse,
+    ActionResponseBody,
+    ActionResponseMetadata,
+    ActionResponseResources,
+    ActionUpdate,
+)
 from zenml.models.v2.core.action_flavor import (
     ActionFlavorResponse,
     ActionFlavorResponseBody,
@@ -318,6 +327,7 @@ from zenml.models.v2.misc.service_connector_type import (
     ResourceTypeModel,
 )
 from zenml.models.v2.misc.server_models import ServerDatabaseType, ServerModel
+from zenml.models.v2.misc.full_stack import FullStackRequest
 from zenml.models.v2.core.trigger import (
     TriggerRequest,
     TriggerFilter,
@@ -377,10 +387,16 @@ from zenml.models.v2.core.server_settings import (
     ServerSettingsResponseMetadata,
     ServerSettingsUpdate,
 )
+from zenml.models.v2.misc.stack_deployment import (
+    DeployedStack,
+    StackDeploymentConfig,
+    StackDeploymentInfo,
+)
 
 # ----------------------------- Forward References -----------------------------
 
 # V2
+ActionResponseResources.model_rebuild()
 APIKeyResponseBody.model_rebuild()
 ArtifactVersionRequest.model_rebuild()
 ArtifactVersionResponseBody.model_rebuild()
@@ -395,6 +411,7 @@ EventSourceResponseMetadata.model_rebuild()
 EventSourceResponseResources.model_rebuild()
 FlavorResponseBody.model_rebuild()
 FlavorResponseMetadata.model_rebuild()
+FullStackRequest.model_rebuild()
 LazyArtifactVersionResponse.model_rebuild()
 LazyRunMetadataResponse.model_rebuild()
 ModelResponseBody.model_rebuild()
@@ -470,6 +487,13 @@ __all__ = [
     "UUIDFilter",
     "Page",
     # V2 Core
+    "ActionFilter",
+    "ActionRequest",
+    "ActionResponse",
+    "ActionResponseBody",
+    "ActionResponseMetadata",
+    "ActionResponseResources",
+    "ActionUpdate",
     "ActionFlavorResponse",
     "ActionFlavorResponseBody",
     "ActionFlavorResponseMetadata",
@@ -684,11 +708,13 @@ __all__ = [
     "WorkspaceResponseMetadata",
     # V2 Misc
     "AuthenticationMethodModel",
+    "DeployedStack",
     "ServiceConnectorResourcesModel",
     "ServiceConnectorTypeModel",
     "ServiceConnectorTypedResourcesModel",
     "ServiceConnectorRequirements",
     "ResourceTypeModel",
+    "FullStackRequest",
     "UserAuthModel",
     "ExternalUserModel",
     "BuildItem",
@@ -701,6 +727,8 @@ __all__ = [
     "ServerModel",
     "ServerDatabaseType",
     "ServerDeploymentType",
+    "StackDeploymentConfig",
+    "StackDeploymentInfo",
     "OAuthDeviceAuthorizationRequest",
     "OAuthDeviceAuthorizationResponse",
     "OAuthDeviceTokenRequest",

zenml/models/v2/core/action.py

@@ -0,0 +1,276 @@
+#  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.
+"""Collection of all models concerning actions."""
+
+import copy
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Optional,
+    TypeVar,
+)
+from uuid import UUID
+
+from pydantic import Field
+
+from zenml.constants import STR_FIELD_MAX_LENGTH
+from zenml.enums import PluginSubType
+from zenml.models.v2.base.base import BaseUpdate
+from zenml.models.v2.base.scoped import (
+    WorkspaceScopedFilter,
+    WorkspaceScopedRequest,
+    WorkspaceScopedResponse,
+    WorkspaceScopedResponseBody,
+    WorkspaceScopedResponseMetadata,
+    WorkspaceScopedResponseResources,
+)
+from zenml.models.v2.core.user import UserResponse
+
+if TYPE_CHECKING:
+    from zenml.zen_stores.schemas import BaseSchema
+
+    AnySchema = TypeVar("AnySchema", bound=BaseSchema)
+
+
+# ------------------ Request Model ------------------
+
+
+class ActionRequest(WorkspaceScopedRequest):
+    """Model for creating a new action."""
+
+    name: str = Field(
+        title="The name of the action.", max_length=STR_FIELD_MAX_LENGTH
+    )
+    description: str = Field(
+        default="",
+        title="The description of the action",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    flavor: str = Field(
+        title="The flavor of the action.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    plugin_subtype: PluginSubType = Field(
+        title="The subtype of the action.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    configuration: Dict[str, Any] = Field(
+        title="The configuration for the action.",
+    )
+    service_account_id: UUID = Field(
+        title="The service account that is used to execute the action.",
+    )
+    auth_window: Optional[int] = Field(
+        default=None,
+        title="The time window in minutes for which the service account is "
+        "authorized to execute the action. Set this to 0 to authorize the "
+        "service account indefinitely (not recommended). If not set, a "
+        "default value defined for each individual action type is used.",
+    )
+
+
+# ------------------ Update Model ------------------
+
+
+class ActionUpdate(BaseUpdate):
+    """Update model for actions."""
+
+    name: Optional[str] = Field(
+        default=None,
+        title="The new name for the action.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    description: Optional[str] = Field(
+        default=None,
+        title="The new description for the action.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    configuration: Optional[Dict[str, Any]] = Field(
+        default=None,
+        title="The configuration for the action.",
+    )
+    service_account_id: Optional[UUID] = Field(
+        default=None,
+        title="The service account that is used to execute the action.",
+    )
+    auth_window: Optional[int] = Field(
+        default=None,
+        title="The time window in minutes for which the service account is "
+        "authorized to execute the action. Set this to 0 to authorize the "
+        "service account indefinitely (not recommended). If not set, a "
+        "default value defined for each individual action type is used.",
+    )
+
+    @classmethod
+    def from_response(cls, response: "ActionResponse") -> "ActionUpdate":
+        """Create an update model from a response model.
+
+        Args:
+            response: The response model to create the update model from.
+
+        Returns:
+            The update model.
+        """
+        return ActionUpdate(
+            configuration=copy.deepcopy(response.configuration),
+        )
+
+
+# ------------------ Response Model ------------------
+
+
+class ActionResponseBody(WorkspaceScopedResponseBody):
+    """Response body for actions."""
+
+    flavor: str = Field(
+        title="The flavor of the action.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    plugin_subtype: PluginSubType = Field(
+        title="The subtype of the action.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+
+
+class ActionResponseMetadata(WorkspaceScopedResponseMetadata):
+    """Response metadata for actions."""
+
+    description: str = Field(
+        default="",
+        title="The description of the action.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    configuration: Dict[str, Any] = Field(
+        title="The configuration for the action.",
+    )
+    auth_window: int = Field(
+        title="The time window in minutes for which the service account is "
+        "authorized to execute the action."
+    )
+
+
+class ActionResponseResources(WorkspaceScopedResponseResources):
+    """Class for all resource models associated with the action entity."""
+
+    service_account: UserResponse = Field(
+        title="The service account that is used to execute the action.",
+    )
+
+
+class ActionResponse(
+    WorkspaceScopedResponse[
+        ActionResponseBody, ActionResponseMetadata, ActionResponseResources
+    ]
+):
+    """Response model for actions."""
+
+    name: str = Field(
+        title="The name of the action.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+
+    def get_hydrated_version(self) -> "ActionResponse":
+        """Get the hydrated version of this action.
+
+        Returns:
+            An instance of the same entity with the metadata field attached.
+        """
+        from zenml.client import Client
+
+        return Client().zen_store.get_action(self.id)
+
+    # Body and metadata properties
+    @property
+    def flavor(self) -> str:
+        """The `flavor` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().flavor
+
+    @property
+    def plugin_subtype(self) -> PluginSubType:
+        """The `plugin_subtype` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().plugin_subtype
+
+    @property
+    def description(self) -> str:
+        """The `description` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().description
+
+    @property
+    def auth_window(self) -> int:
+        """The `auth_window` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().auth_window
+
+    @property
+    def configuration(self) -> Dict[str, Any]:
+        """The `configuration` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().configuration
+
+    def set_configuration(self, configuration: Dict[str, Any]) -> None:
+        """Set the `configuration` property.
+
+        Args:
+            configuration: The value to set.
+        """
+        self.get_metadata().configuration = configuration
+
+    # Resource properties
+    @property
+    def service_account(self) -> "UserResponse":
+        """The `service_account` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_resources().service_account
+
+
+# ------------------ Filter Model ------------------
+
+
+class ActionFilter(WorkspaceScopedFilter):
+    """Model to enable advanced filtering of all actions."""
+
+    name: Optional[str] = Field(
+        default=None,
+        description="Name of the action.",
+    )
+    flavor: Optional[str] = Field(
+        default=None,
+        title="The flavor of the action.",
+    )
+    plugin_subtype: Optional[str] = Field(
+        default=None,
+        title="The subtype of the action.",
+    )

zenml/models/v2/core/component.py

@@ -238,6 +238,24 @@ class ComponentResponse(
         max_length=STR_FIELD_MAX_LENGTH,
     )
 
+    def get_analytics_metadata(self) -> Dict[str, Any]:
+        """Add the component labels to analytics metadata.
+
+        Returns:
+            Dict of analytics metadata.
+        """
+        metadata = super().get_analytics_metadata()
+
+        if self.labels is not None:
+            metadata.update(
+                {
+                    label[6:]: value
+                    for label, value in self.labels.items()
+                    if label.startswith("zenml:")
+                }
+            )
+        return metadata
+
     def get_hydrated_version(self) -> "ComponentResponse":
         """Get the hydrated version of this component.
 

zenml/models/v2/core/model.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Models representing models."""
 
-from functools import partial
 from typing import TYPE_CHECKING, ClassVar, List, Optional, Union
 from uuid import UUID
 
@@ -303,7 +302,7 @@ class ModelResponse(
 
         client = Client()
         model_versions = depaginate(
-            partial(client.list_model_versions, model_name_or_id=self.id)
+            client.list_model_versions, model_name_or_id=self.id
         )
         return [
             mv.to_model_class(suppress_class_validation_warnings=True)

zenml/models/v2/core/service_connector.py

@@ -493,6 +493,23 @@ class ServiceConnectorResponse(
         max_length=STR_FIELD_MAX_LENGTH,
     )
 
+    def get_analytics_metadata(self) -> Dict[str, Any]:
+        """Add the service connector labels to analytics metadata.
+
+        Returns:
+            Dict of analytics metadata.
+        """
+        metadata = super().get_analytics_metadata()
+
+        metadata.update(
+            {
+                label[6:]: value
+                for label, value in self.labels.items()
+                if label.startswith("zenml:")
+            }
+        )
+        return metadata
+
     def get_hydrated_version(self) -> "ServiceConnectorResponse":
         """Get the hydrated version of this service connector.
 

zenml/models/v2/core/stack.py

@@ -36,6 +36,7 @@ from zenml.models.v2.core.component import ComponentResponse
 if TYPE_CHECKING:
     from sqlalchemy.sql.elements import ColumnElement
 
+
 # ------------------ Request Model ------------------
 
 
@@ -59,6 +60,10 @@ class StackRequest(WorkspaceScopedRequest):
         title="A mapping of stack component types to the actual"
         "instances of components of this type.",
     )
+    labels: Optional[Dict[str, Any]] = Field(
+        default=None,
+        title="The stack labels.",
+    )
 
     @property
     def is_valid(self) -> bool:
@@ -109,6 +114,10 @@ class StackUpdate(BaseUpdate):
         "instances of components of this type.",
         default=None,
     )
+    labels: Optional[Dict[str, Any]] = Field(
+        default=None,
+        title="The stack labels.",
+    )
 
 
 # ------------------ Response Model ------------------
@@ -134,6 +143,10 @@ class StackResponseMetadata(WorkspaceScopedResponseMetadata):
         default=None,
         title="The path to the stack spec used for mlstacks deployments.",
     )
+    labels: Optional[Dict[str, Any]] = Field(
+        default=None,
+        title="The stack labels.",
+    )
 
 
 class StackResponseResources(WorkspaceScopedResponseResources):
@@ -214,6 +227,15 @@ class StackResponse(
         """
         metadata = super().get_analytics_metadata()
         metadata.update({ct: c[0].flavor for ct, c in self.components.items()})
+
+        if self.labels is not None:
+            metadata.update(
+                {
+                    label[6:]: value
+                    for label, value in self.labels.items()
+                    if label.startswith("zenml:")
+                }
+            )
         return metadata
 
     @property
@@ -243,6 +265,15 @@ class StackResponse(
         """
         return self.get_metadata().components
 
+    @property
+    def labels(self) -> Optional[Dict[str, Any]]:
+        """The `labels` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().labels
+
 
 # ------------------ Filter Model ------------------