zenml-nightly 0.58.2.dev20240626__py3-none-any.whl → 0.62.0.dev20240726__py3-none-any.whl

zenml/integrations/skypilot_gcp/__init__.py

@@ -31,7 +31,7 @@ class SkypilotGCPIntegration(Integration):
     """Definition of Skypilot (GCP) Integration for ZenML."""
 
     NAME = SKYPILOT_GCP
-    REQUIREMENTS = ["skypilot[gcp]<=0.5.0"]
+    REQUIREMENTS = ["skypilot[gcp]~=0.6.0"]
     APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod

zenml/integrations/skypilot_lambda/__init__.py

@@ -31,7 +31,7 @@ class SkypilotLambdaIntegration(Integration):
     """Definition of Skypilot Lambda Integration for ZenML."""
 
     NAME = SKYPILOT_LAMBDA
-    REQUIREMENTS = ["skypilot[lambda]<=0.5.0"]
+    REQUIREMENTS = ["skypilot[lambda]~=0.6.0"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/skypilot_lambda/flavors/skypilot_orchestrator_lambda_vm_flavor.py

@@ -40,7 +40,7 @@ class SkypilotLambdaOrchestratorSettings(SkypilotBaseOrchestratorSettings):
 
     _UNSUPPORTED_FEATURES = {
         "use_spot": "Spot instances not supported for Lambda orchestrator.",
-        "spot_recovery": "Spot recovery not supported for Lambda orchestrator.",
+        "job_recovery": "Job recovery not supported for Lambda orchestrator.",
         "image_id": "Custom image IDs not supported for Lambda orchestrator.",
         # Add other unsupported features as needed
     }

zenml/integrations/slack/__init__.py

@@ -31,6 +31,7 @@ class SlackIntegration(Integration):
 
     NAME = SLACK
     REQUIREMENTS = ["slack-sdk>=3.16.1", "aiohttp>=3.8.1"]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["aiohttp"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/tekton/__init__.py

@@ -32,6 +32,7 @@ class TektonIntegration(Integration):
 
     NAME = TEKTON
     REQUIREMENTS = ["kfp>=2.6.0", "kfp-kubernetes>=1.1.0"]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["kfp"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/tensorboard/__init__.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Initialization for TensorBoard integration."""
 
-import sys
 from typing import List, Optional
 from zenml.integrations.constants import TENSORBOARD
 from zenml.integrations.integration import Integration

zenml/integrations/tensorflow/__init__.py

@@ -18,6 +18,9 @@ import sys
 from typing import List, Optional
 from zenml.integrations.constants import TENSORFLOW
 from zenml.integrations.integration import Integration
+from zenml.logger import get_logger
+
+logger = get_logger(__name__)
 
 
 class TensorflowIntegration(Integration):
@@ -25,20 +28,27 @@ class TensorflowIntegration(Integration):
 
     NAME = TENSORFLOW
     REQUIREMENTS = []
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["typing-extensions"]
 
     @classmethod
     def activate(cls) -> None:
         """Activates the integration."""
         # need to import this explicitly to load the Tensorflow file IO support
         # for S3 and other file systems
-        if (
-            not platform.system() == "Darwin"
-            or not platform.machine() == "arm64"
-        ):
+        if not platform.system() == "Darwin" or not platform.machine() == "arm64":
             import tensorflow_io  # type: ignore
 
         from zenml.integrations.tensorflow import materializers  # noqa
 
+        if sys.version_info.minor == 8:
+            logger.warning(
+                "Python 3.8 with TensorFlow is not fully "
+                "compatible with Pydantic 2 requirements. "
+                "Consider upgrading to a higher Python "
+                "version if you would like to use the "
+                "Tensorflow integration."
+            )
+
     @classmethod
     def get_requirements(cls, target_os: Optional[str] = None) -> List[str]:
         """Defines platform specific requirements for the integration.
@@ -52,13 +62,15 @@ class TensorflowIntegration(Integration):
         target_os = target_os or platform.system()
         if target_os == "Darwin" and platform.machine() == "arm64":
             requirements = [
-                f"tensorflow-macos>=2.12,<=2.15",
+                "tensorflow-macos>=2.12,<=2.15",
             ]
         else:
             requirements = [
-                f"tensorflow>=2.12,<=2.15",
+                "tensorflow>=2.12,<=2.15",
                 "tensorflow_io>=0.24.0",
             ]
+        if sys.version_info.minor == 8:
+            requirements.append("typing-extensions>=4.6.1")
         return requirements
 
 

zenml/integrations/wandb/__init__.py

@@ -31,6 +31,7 @@ class WandbIntegration(Integration):
 
     NAME = WANDB
     REQUIREMENTS = ["wandb>=0.12.12", "Pillow>=9.1.0"]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["Pillow"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/logging/step_logging.py

@@ -240,8 +240,6 @@ class StepLogsStorage:
 
         # Immutable filesystems state
         self.last_merge_time = time.time()
-        self.log_files_not_merged: List[str] = []
-        self.next_merged_file_name: str = self._get_timestamped_filename()
 
     @property
     def artifact_store(self) -> "BaseArtifactStore":
@@ -279,13 +277,16 @@ class StepLogsStorage:
             or time.time() - self.last_save_time >= self.time_interval
         )
 
-    def _get_timestamped_filename(self) -> str:
+    def _get_timestamped_filename(self, suffix: str = "") -> str:
         """Returns a timestamped filename.
 
+        Args:
+            suffix: optional suffix for the file name
+
         Returns:
             The timestamped filename.
         """
-        return f"{time.time()}{LOGS_EXTENSION}"
+        return f"{time.time()}{suffix}{LOGS_EXTENSION}"
 
     def save_to_file(self, force: bool = False) -> None:
         """Method to save the buffer to the given URI.
@@ -302,12 +303,7 @@ class StepLogsStorage:
             try:
                 if self.buffer:
                     if self.artifact_store.config.IS_IMMUTABLE_FILESYSTEM:
-                        if not self.log_files_not_merged:
-                            self.next_merged_file_name = (
-                                self._get_timestamped_filename()
-                            )
                         _logs_uri = self._get_timestamped_filename()
-                        self.log_files_not_merged.append(_logs_uri)
                         with self.artifact_store.open(
                             os.path.join(
                                 self.logs_uri,
@@ -346,42 +342,40 @@ class StepLogsStorage:
             and time.time() - self.last_merge_time > self.merge_files_interval
         ):
             try:
-                self.merge_log_files(
-                    self.next_merged_file_name, self.log_files_not_merged
-                )
+                self.merge_log_files()
             except (OSError, IOError) as e:
                 logger.error(f"Error while trying to roll up logs: {e}")
-            else:
-                self.log_files_not_merged = []
             finally:
                 self.last_merge_time = time.time()
 
-    def merge_log_files(
-        self,
-        file_name: Optional[str] = None,
-        files: Optional[List[str]] = None,
-    ) -> None:
+    def merge_log_files(self, merge_all_files: bool = False) -> None:
         """Merges all log files into one in the given URI.
 
         Called on the logging context exit.
 
         Args:
-            file_name: The name of the merged log file.
-            files: The list of log files to merge.
+            merge_all_files: whether to merge all files or only raw files
         """
         if self.artifact_store.config.IS_IMMUTABLE_FILESYSTEM:
-            files_ = files or self.artifact_store.listdir(self.logs_uri)
-            file_name_ = file_name or self._get_timestamped_filename()
+            merged_file_suffix = "_merged"
+            files_ = self.artifact_store.listdir(self.logs_uri)
+            if not merge_all_files:
+                # already merged files will not be merged again
+                files_ = [f for f in files_ if merged_file_suffix not in f]
+            file_name_ = self._get_timestamped_filename(
+                suffix=merged_file_suffix
+            )
             if len(files_) > 1:
                 files_.sort()
                 logger.debug("Log files count: %s", len(files_))
 
-                try:
-                    # dump all logs to a local file first
-                    with self.artifact_store.open(
-                        os.path.join(self.logs_uri, file_name_), "w"
-                    ) as merged_file:
-                        for file in files_:
+                missing_files = set()
+                # dump all logs to a local file first
+                with self.artifact_store.open(
+                    os.path.join(self.logs_uri, file_name_), "w"
+                ) as merged_file:
+                    for file in files_:
+                        try:
                             merged_file.write(
                                 str(
                                     _load_file_from_artifact_store(
@@ -391,11 +385,12 @@ class StepLogsStorage:
                                     )
                                 )
                             )
-                except Exception as e:
-                    logger.warning(f"Failed to merge log files. {e}")
-                else:
-                    # clean up left over files
-                    for file in files_:
+                        except DoesNotExistException:
+                            missing_files.add(file)
+
+                # clean up left over files
+                for file in files_:
+                    if file not in missing_files:
                         self.artifact_store.remove(
                             os.path.join(self.logs_uri, str(file))
                         )
@@ -452,7 +447,6 @@ class StepLogsStorageContext:
         Restores the `write` method of both stderr and stdout.
         """
         self.storage.save_to_file(force=True)
-        self.storage.merge_log_files()
 
         setattr(sys.stdout, "write", self.stdout_write)
         setattr(sys.stdout, "flush", self.stdout_flush)
@@ -462,6 +456,11 @@ class StepLogsStorageContext:
 
         redirected.set(False)
 
+        try:
+            self.storage.merge_log_files(merge_all_files=True)
+        except (OSError, IOError) as e:
+            logger.warning(f"Step logs roll-up failed: {e}")
+
     def _wrap_write(self, method: Callable[..., Any]) -> Callable[..., Any]:
         """Wrapper function that utilizes the storage object to store logs.
 

zenml/materializers/built_in_materializer.py

@@ -80,7 +80,7 @@ class BuiltInMaterializer(BaseMaterializer):
             The data read.
         """
         contents = yaml_utils.read_json(self.data_path)
-        if type(contents) != data_type:
+        if type(contents) is not data_type:
             # TODO [ENG-142]: Raise error or try to coerce
             logger.debug(
                 f"Contents {contents} was type {type(contents)} but expected "

zenml/materializers/cloudpickle_materializer.py

@@ -94,7 +94,7 @@ class CloudpickleMaterializer(BaseMaterializer):
         """
         # Log a warning if this materializer was not explicitly specified for
         # the given data type.
-        if type(self) == CloudpickleMaterializer:
+        if type(self) is CloudpickleMaterializer:
             logger.warning(
                 f"No materializer is registered for type `{type(data)}`, so "
                 "the default Pickle materializer was used. Pickle is not "

zenml/model/model.py

@@ -518,7 +518,7 @@ class Model(BaseModel):
             and not suppress_class_validation_warnings
         ):
             logger.info(
-                f"`version` `{version}` matches one of the possible "
+                f"Version `{version}` matches one of the possible "
                 "`ModelStages` and will be fetched using stage."
             )
         if str(version).isnumeric() and not suppress_class_validation_warnings:

zenml/models/__init__.py

@@ -327,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,
@@ -386,6 +387,11 @@ from zenml.models.v2.core.server_settings import (
     ServerSettingsResponseMetadata,
     ServerSettingsUpdate,
 )
+from zenml.models.v2.misc.stack_deployment import (
+    DeployedStack,
+    StackDeploymentConfig,
+    StackDeploymentInfo,
+)
 
 # ----------------------------- Forward References -----------------------------
 
@@ -405,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()
@@ -701,11 +708,13 @@ __all__ = [
     "WorkspaceResponseMetadata",
     # V2 Misc
     "AuthenticationMethodModel",
+    "DeployedStack",
     "ServiceConnectorResourcesModel",
     "ServiceConnectorTypeModel",
     "ServiceConnectorTypedResourcesModel",
     "ServiceConnectorRequirements",
     "ResourceTypeModel",
+    "FullStackRequest",
     "UserAuthModel",
     "ExternalUserModel",
     "BuildItem",
@@ -718,6 +727,8 @@ __all__ = [
     "ServerModel",
     "ServerDatabaseType",
     "ServerDeploymentType",
+    "StackDeploymentConfig",
+    "StackDeploymentInfo",
     "OAuthDeviceAuthorizationRequest",
     "OAuthDeviceAuthorizationResponse",
     "OAuthDeviceTokenRequest",

zenml/models/v2/core/component.py

@@ -191,6 +191,17 @@ class ComponentResponseBody(WorkspaceScopedResponseBody):
         title="The flavor of the stack component.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
+    integration: Optional[str] = Field(
+        default=None,
+        title="The name of the integration that the component's flavor "
+        "belongs to.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    logo_url: Optional[str] = Field(
+        default=None,
+        title="Optionally, a url pointing to a png,"
+        "svg or jpg can be attached.",
+    )
 
 
 class ComponentResponseMetadata(WorkspaceScopedResponseMetadata):
@@ -238,6 +249,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.
 
@@ -267,6 +296,24 @@ class ComponentResponse(
         """
         return self.get_body().flavor
 
+    @property
+    def integration(self) -> Optional[str]:
+        """The `integration` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().integration
+
+    @property
+    def logo_url(self) -> Optional[str]:
+        """The `logo_url` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().logo_url
+
     @property
     def configuration(self) -> Dict[str, Any]:
         """The `configuration` property.

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/server_settings.py

@@ -15,8 +15,6 @@
 
 from datetime import datetime
 from typing import (
-    Any,
-    Dict,
     Optional,
 )
 from uuid import UUID
@@ -57,10 +55,6 @@ class ServerSettingsUpdate(BaseZenModel):
         default=None,
         title="Whether to display notifications about ZenML updates in the dashboard.",
     )
-    onboarding_state: Optional[Dict[str, Any]] = Field(
-        default=None,
-        title="The server's onboarding state.",
-    )
 
 
 # ------------------ Response Model ------------------
@@ -96,11 +90,6 @@ class ServerSettingsResponseBody(BaseResponseBody):
 class ServerSettingsResponseMetadata(BaseResponseMetadata):
     """Response metadata for server settings."""
 
-    onboarding_state: Dict[str, Any] = Field(
-        default={},
-        title="The server's onboarding state.",
-    )
-
 
 class ServerSettingsResponseResources(BaseResponseResources):
     """Response resources for server settings."""
@@ -199,15 +188,6 @@ class ServerSettingsResponse(
         """
         return self.get_body().updated
 
-    @property
-    def onboarding_state(self) -> Dict[str, Any]:
-        """The `onboarding_state` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_metadata().onboarding_state
-
 
 # ------------------ Filter Model ------------------
 

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 ------------------