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

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -38,7 +38,6 @@ from kubernetes import config as k8s_config
 
 from zenml.config.base_settings import BaseSettings
 from zenml.enums import StackComponentType
-from zenml.environment import Environment
 from zenml.integrations.kubernetes.flavors.kubernetes_orchestrator_flavor import (
     KubernetesOrchestratorConfig,
     KubernetesOrchestratorSettings,
@@ -335,19 +334,8 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                 environment.
 
         Raises:
-            RuntimeError: If trying to run from a Jupyter notebook.
+            RuntimeError: If the Kubernetes orchestrator is not configured.
         """
-        # First check whether the code is running in a notebook.
-        if Environment.in_notebook():
-            raise RuntimeError(
-                "The Kubernetes orchestrator cannot run pipelines in a notebook "
-                "environment. The reason is that it is non-trivial to create "
-                "a Docker image of a notebook. Please consider refactoring "
-                "your notebook cells into separate scripts in a Python module "
-                "and run the code outside of a notebook when using this "
-                "orchestrator."
-            )
-
         for step_name, step in deployment.step_configurations.items():
             if self.requires_resources_in_orchestration_environment(step):
                 logger.warning(

zenml/integrations/kubernetes/orchestrators/manifest_utils.py

@@ -139,24 +139,42 @@ def build_pod_manifest(
         ],
         security_context=security_context,
     )
+    image_pull_secrets = []
+    if pod_settings:
+        image_pull_secrets = [
+            k8s_client.V1LocalObjectReference(name=name)
+            for name in pod_settings.image_pull_secrets
+        ]
 
     pod_spec = k8s_client.V1PodSpec(
         containers=[container_spec],
         restart_policy="Never",
+        image_pull_secrets=image_pull_secrets,
     )
 
     if service_account_name is not None:
         pod_spec.service_account_name = service_account_name
 
+    labels = {}
+
     if pod_settings:
         add_pod_settings(pod_spec, pod_settings)
 
-    pod_metadata = k8s_client.V1ObjectMeta(
-        name=pod_name,
-        labels={
+        # Add pod_settings.labels to the labels
+        if pod_settings.labels:
+            labels.update(pod_settings.labels)
+
+    # Add run_name and pipeline_name to the labels
+    labels.update(
+        {
             "run": run_name,
             "pipeline": pipeline_name,
-        },
+        }
+    )
+
+    pod_metadata = k8s_client.V1ObjectMeta(
+        name=pod_name,
+        labels=labels,
     )
 
     if pod_settings and pod_settings.annotations:

zenml/integrations/kubernetes/pod_settings.py

@@ -33,6 +33,8 @@ class KubernetesPodSettings(BaseSettings):
         volumes: Volumes to mount in the pod.
         volume_mounts: Volume mounts to apply to the pod containers.
         host_ipc: Whether to enable host IPC for the pod.
+        image_pull_secrets: Image pull secrets to use for the pod.
+        labels: Labels to apply to the pod.
     """
 
     node_selectors: Dict[str, str] = {}
@@ -43,6 +45,8 @@ class KubernetesPodSettings(BaseSettings):
     volumes: List[Dict[str, Any]] = []
     volume_mounts: List[Dict[str, Any]] = []
     host_ipc: bool = False
+    image_pull_secrets: List[str] = []
+    labels: Dict[str, str] = {}
 
     @field_validator("volumes", mode="before")
     @classmethod

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