zenml-nightly 0.61.0.dev20240712__py3-none-any.whl → 0.62.0.dev20240727__py3-none-any.whl

zenml/integrations/lightgbm/__init__.py

@@ -22,6 +22,7 @@ class LightGBMIntegration(Integration):
 
     NAME = LIGHTGBM
     REQUIREMENTS = ["lightgbm>=1.0.0"]
+    APT_PACKAGES = ["libgomp1"]
 
     @classmethod
     def activate(cls) -> None:

zenml/integrations/mlflow/__init__.py

@@ -33,7 +33,7 @@ class MlflowIntegration(Integration):
     NAME = MLFLOW
 
     REQUIREMENTS = [
-        "mlflow>=2.1.1,<=2.14.1",
+        "mlflow>=2.1.1,<=2.14.2",
         "mlserver>=1.3.3",
         "mlserver-mlflow>=1.3.3",
         # TODO: remove this requirement once rapidjson is fixed

zenml/integrations/mlflow/model_registries/mlflow_model_registry.py

@@ -371,7 +371,7 @@ class MLFlowModelRegistry(BaseModelRegistry):
             self.get_model(name=name)
         except KeyError:
             logger.info(
-                f"No registered model with name {name} found. Creating a new"
+                f"No registered model with name {name} found. Creating a new "
                 "registered model."
             )
             self.register_model(
@@ -615,7 +615,11 @@ class MLFlowModelRegistry(BaseModelRegistry):
         for mlflow_model_version in mlflow_model_versions:
             # check if given MlFlow model version matches the given request
             # before casting it
-            if stage and not mlflow_model_version.current_stage == str(stage):
+            if (
+                stage
+                and not ModelVersionStage(mlflow_model_version.current_stage)
+                == stage
+            ):
                 continue
             if created_after and not (
                 mlflow_model_version.creation_timestamp

zenml/integrations/mlflow/services/mlflow_deployment.py

@@ -260,7 +260,7 @@ class MLFlowDeploymentService(LocalDaemonService, BaseDeploymentService):
             )
 
         if self.endpoint.prediction_url is not None:
-            if type(request) == pd.DataFrame:
+            if type(request) is pd.DataFrame:
                 response = requests.post(  # nosec
                     self.endpoint.prediction_url,
                     json={"instances": request.to_dict("records")},

zenml/integrations/skypilot_azure/__init__.py

@@ -19,13 +19,11 @@ orchestrator for a remote orchestration of ZenML pipelines on VMs.
 from typing import List, Type
 
 from zenml.integrations.constants import (
-
     SKYPILOT_AZURE,
 )
 from zenml.integrations.integration import Integration
 from zenml.stack import Flavor
 
-
 SKYPILOT_AZURE_ORCHESTRATOR_FLAVOR = "vm_azure"
 
 
@@ -33,7 +31,7 @@ class SkypilotAzureIntegration(Integration):
     """Definition of Skypilot (Azure) Integration for ZenML."""
 
     NAME = SKYPILOT_AZURE
-    REQUIREMENTS = ["skypilot[azure]~=0.6.0"]
+    REQUIREMENTS = ["skypilot-nightly[azure]==1.0.0.dev20240716"]
     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.6.0"]
+    REQUIREMENTS = ["skypilot[lambda]~=0.6.0"]
 
     @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/v2/core/code_repository.py

@@ -18,7 +18,7 @@ from uuid import UUID
 
 from pydantic import Field
 
-from zenml.config.source import Source
+from zenml.config.source import Source, SourceWithValidator
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
 from zenml.models.v2.base.base import BaseUpdate
 from zenml.models.v2.base.scoped import (
@@ -71,7 +71,7 @@ class CodeRepositoryUpdate(BaseUpdate):
         description="Configuration for the code repository.",
         default=None,
     )
-    source: Optional[Source] = Field(
+    source: Optional[SourceWithValidator] = Field(
         description="The code repository source.", default=None
     )
     logo_url: Optional[str] = Field(

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):
@@ -285,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/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/misc/full_stack.py

@@ -21,6 +21,7 @@ from pydantic import BaseModel, Field, model_validator
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.enums import StackComponentType
 from zenml.models.v2.base.base import BaseRequest
+from zenml.models.v2.core.component import ComponentResponse
 
 
 class ServiceConnectorInfo(BaseModel):
@@ -95,3 +96,34 @@ class FullStackRequest(BaseRequest):
                             "the position in the list of service connectors."
                         )
         return self
+
+
+class ResourcesInfo(BaseModel):
+    """Information about the resources needed for CLI and UI."""
+
+    flavor: str
+    flavor_display_name: str
+    required_configuration: Dict[str, str] = {}
+    use_resource_value_as_fixed_config: bool = False
+
+    accessible_by_service_connector: List[str]
+    connected_through_service_connector: List[ComponentResponse]
+
+    @model_validator(mode="after")
+    def _validate_resource_info(self) -> "ResourcesInfo":
+        if (
+            self.use_resource_value_as_fixed_config
+            and len(self.required_configuration) > 1
+        ):
+            raise ValueError(
+                "Cannot use resource value as fixed config if more than one required configuration key is provided."
+            )
+        return self
+
+
+class ServiceConnectorResourcesInfo(BaseModel):
+    """Information about the service connector resources needed for CLI and UI."""
+
+    connector_type: str
+
+    components_resources_info: Dict[StackComponentType, List[ResourcesInfo]]

zenml/models/v2/misc/stack_deployment.py

@@ -55,6 +55,11 @@ class StackDeploymentInfo(BaseModel):
         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):

zenml/new/pipelines/run_utils.py

@@ -250,7 +250,7 @@ def _validate_new_version_requests(
     if not is_cloud_model:
         logger.info(
             "Models can be viewed in the dashboard using ZenML Pro. Sign up "
-            "for a free trial at https://www.zenml.io/cloud/"
+            "for a free trial at https://www.zenml.io/pro/"
         )
 
 

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

@@ -291,6 +291,7 @@ class StepLauncher:
                                 logger.error(
                                     f"Failed to run step `{self._step_name}` after {max_retries} retries. Exiting."
                                 )
+                                logger.exception(e)
                                 publish_utils.publish_failed_step_run(
                                     step_run_response.id
                                 )

zenml/orchestrators/wheeled_orchestrator.py

@@ -0,0 +1,147 @@
+#  Copyright (c) ZenML GmbH 2021. 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.
+"""Wheeled orchestrator class."""
+
+import os
+import re
+import subprocess
+import tempfile
+from abc import ABC
+
+from zenml import __version__
+from zenml.io import fileio
+from zenml.logger import get_logger
+from zenml.orchestrators import BaseOrchestrator
+from zenml.utils.io_utils import copy_dir
+from zenml.utils.source_utils import get_source_root
+
+logger = get_logger(__name__)
+
+DEFAULT_PACKAGE_NAME = "zenmlproject"
+
+
+class WheeledOrchestrator(BaseOrchestrator, ABC):
+    """Base class for wheeled orchestrators."""
+
+    package_name = DEFAULT_PACKAGE_NAME
+    package_version = __version__
+
+    def copy_repository_to_temp_dir_and_add_setup_py(self) -> str:
+        """Copy the repository to a temporary directory and add a setup.py file.
+
+        Returns:
+            Path to the temporary directory containing the copied repository.
+        """
+        repo_path = get_source_root()
+
+        self.package_name = f"{DEFAULT_PACKAGE_NAME}_{self.sanitize_name(os.path.basename(repo_path))}"
+
+        # Create a temporary folder
+        temp_dir = tempfile.mkdtemp(prefix="zenml-temp-")
+
+        # Create a folder within the temporary directory
+        temp_repo_path = os.path.join(temp_dir, self.package_name)
+        fileio.mkdir(temp_repo_path)
+
+        # Copy the repository to the temporary directory
+        copy_dir(repo_path, temp_repo_path)
+
+        # Create init file in the copied directory
+        init_file_path = os.path.join(temp_repo_path, "__init__.py")
+        with fileio.open(init_file_path, "w") as f:
+            f.write("")
+
+        # Create a setup.py file
+        setup_py_content = f"""
+from setuptools import setup, find_packages
+
+setup(
+    name="{self.package_name}",
+    version="{self.package_version}",
+    packages=find_packages(),
+)
+"""
+        setup_py_path = os.path.join(temp_dir, "setup.py")
+        with fileio.open(setup_py_path, "w") as f:
+            f.write(setup_py_content)
+
+        return temp_dir
+
+    def create_wheel(self, temp_dir: str) -> str:
+        """Create a wheel for the package in the given temporary directory.
+
+        Args:
+            temp_dir (str): Path to the temporary directory containing the package.
+
+        Raises:
+            RuntimeError: If the wheel file could not be created.
+
+        Returns:
+            str: Path to the created wheel file.
+        """
+        # Change to the temporary directory
+        original_dir = os.getcwd()
+        os.chdir(temp_dir)
+
+        try:
+            # Run the `pip wheel` command to create the wheel
+            result = subprocess.run(
+                ["pip", "wheel", "."], check=True, capture_output=True
+            )
+            logger.debug(f"Wheel creation stdout: {result.stdout.decode()}")
+            logger.debug(f"Wheel creation stderr: {result.stderr.decode()}")
+
+            # Find the created wheel file
+            wheel_file = next(
+                (
+                    file
+                    for file in os.listdir(temp_dir)
+                    if file.endswith(".whl")
+                ),
+                None,
+            )
+
+            if wheel_file is None:
+                raise RuntimeError("Failed to create wheel file.")
+
+            wheel_path = os.path.join(temp_dir, wheel_file)
+
+            # Verify the wheel file is a valid zip file
+            import zipfile
+
+            if not zipfile.is_zipfile(wheel_path):
+                raise RuntimeError(
+                    f"The file {wheel_path} is not a valid zip file."
+                )
+
+            return wheel_path
+        finally:
+            # Change back to the original directory
+            os.chdir(original_dir)
+
+    def sanitize_name(self, name: str) -> str:
+        """Sanitize the value to be used in a cluster name.
+
+        Args:
+            name: Arbitrary input cluster name.
+
+        Returns:
+            Sanitized cluster name.
+        """
+        name = re.sub(
+            r"[^a-z0-9-]", "-", name.lower()
+        )  # replaces any character that is not a lowercase letter, digit, or hyphen with a hyphen
+        name = re.sub(r"^[-]+", "", name)  # trim leading hyphens
+        name = re.sub(r"[-]+$", "", name)  # trim trailing hyphens
+        return name