zenml-nightly 0.61.0.dev20240713__py3-none-any.whl → 0.62.0.dev20240717__py3-none-any.whl

zenml/integrations/kubernetes/pod_settings.py

@@ -33,6 +33,7 @@ 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.
     """
 
     node_selectors: Dict[str, str] = {}
@@ -43,6 +44,7 @@ class KubernetesPodSettings(BaseSettings):
     volumes: List[Dict[str, Any]] = []
     volume_mounts: List[Dict[str, Any]] = []
     host_ipc: bool = False
+    image_pull_secrets: List[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_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/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/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/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/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

zenml/service_connectors/service_connector_utils.py

@@ -0,0 +1,349 @@
+#  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.
+"""Utility methods for Service Connectors."""
+
+from typing import Dict, List, Union
+from uuid import UUID
+
+from zenml.client import Client
+from zenml.enums import StackComponentType
+from zenml.models.v2.core.service_connector import ServiceConnectorRequest
+from zenml.models.v2.misc.full_stack import (
+    ResourcesInfo,
+    ServiceConnectorInfo,
+    ServiceConnectorResourcesInfo,
+)
+from zenml.utils.pagination_utils import depaginate
+
+
+def _prepare_resource_info(
+    connector_details: Union[UUID, ServiceConnectorInfo],
+    resource_ids: List[str],
+    stack_component_type: StackComponentType,
+    flavor: str,
+    required_configuration: Dict[str, str],
+    flavor_display_name: str,
+    use_resource_value_as_fixed_config: bool = False,
+) -> ResourcesInfo:
+    existing_components = []
+    if isinstance(connector_details, UUID):
+        existing_components = depaginate(
+            Client().list_stack_components,
+            type=stack_component_type.value,
+            connector_id=connector_details,
+            flavor=flavor,
+        )
+    return ResourcesInfo(
+        flavor=flavor,
+        required_configuration=required_configuration,
+        flavor_display_name=flavor_display_name,
+        use_resource_value_as_fixed_config=use_resource_value_as_fixed_config,
+        accessible_by_service_connector=resource_ids,
+        connected_through_service_connector=existing_components,
+    )
+
+
+def _raise_specific_cloud_exception_if_needed(
+    cloud_provider: str,
+    artifact_stores: List[ResourcesInfo],
+    orchestrators: List[ResourcesInfo],
+    container_registries: List[ResourcesInfo],
+) -> None:
+    AWS_DOCS = (
+        "https://docs.zenml.io/how-to/auth-management/aws-service-connector"
+    )
+    GCP_DOCS = (
+        "https://docs.zenml.io/how-to/auth-management/gcp-service-connector"
+    )
+
+    if not artifact_stores:
+        error_msg = (
+            "We were unable to find any {obj_name} available "
+            "to configured service connector. Please, verify "
+            "that needed permission are granted for the "
+            "service connector.\nDocumentation for the "
+            "{obj_name} configuration can be found at "
+            "{docs}"
+        )
+        if cloud_provider == "aws":
+            raise ValueError(
+                error_msg.format(
+                    obj_name="S3 Bucket", docs=f"{AWS_DOCS}#s3-bucket"
+                )
+            )
+        elif cloud_provider == "gcp":
+            raise ValueError(
+                error_msg.format(
+                    obj_name="GCS Bucket", docs=f"{GCP_DOCS}#gcs-bucket"
+                )
+            )
+        elif cloud_provider == "azure":
+            pass
+    if not orchestrators:
+        error_msg = (
+            "We were unable to find any orchestrator engines "
+            "available to the service connector. Please, verify "
+            "that needed permission are granted for the "
+            "service connector.\nDocumentation for the Generic "
+            "{cloud_name} resource configuration can be found at "
+            "{gen_docs}\n Documentation for the {k8s_name} resource "
+            "configuration can be found at {k8s_docs}"
+        )
+        if cloud_provider == "aws":
+            raise ValueError(
+                error_msg.format(
+                    cloud_name="AWS",
+                    gen_docs=f"{AWS_DOCS}#generic-aws-resource",
+                    k8s_name="Kubernetes",
+                    k8s_docs=f"{AWS_DOCS}#eks-kubernetes-cluster",
+                )
+            )
+
+        elif cloud_provider == "gcp":
+            raise ValueError(
+                error_msg.format(
+                    cloud_name="GCP",
+                    gen_docs=f"{GCP_DOCS}#generic-gcp-resource",
+                    k8s_name="GKE Kubernetes",
+                    k8s_docs=f"{GCP_DOCS}#gke-kubernetes-cluster",
+                )
+            )
+        elif cloud_provider == "azure":
+            pass
+    if not container_registries:
+        error_msg = (
+            "We were unable to find any container registries "
+            "available to the service connector. Please, verify "
+            "that needed permission are granted for the "
+            "service connector.\nDocumentation for the {registry_name} "
+            "container registry resource configuration can "
+            "be found at {docs_link}"
+        )
+        if cloud_provider == "aws":
+            raise ValueError(
+                error_msg.format(
+                    registry_name="ECR",
+                    docs_link=f"{AWS_DOCS}#ecr-container-registry",
+                )
+            )
+        elif cloud_provider == "gcp":
+            raise ValueError(
+                error_msg.format(
+                    registry_name="GCR",
+                    docs_link=f"{GCP_DOCS}#gcr-container-registry",
+                )
+            )
+        elif cloud_provider == "azure":
+            pass
+
+
+def get_resources_options_from_resource_model_for_full_stack(
+    connector_details: Union[UUID, ServiceConnectorInfo],
+) -> ServiceConnectorResourcesInfo:
+    """Get the resource options from the resource model for the full stack.
+
+    Args:
+        connector_details: The service connector details (UUID or Info).
+
+    Returns:
+        All available service connector resource options.
+    """
+    client = Client()
+    zen_store = client.zen_store
+
+    can_generate_long_tokens = False
+    if isinstance(connector_details, UUID):
+        resource_model = zen_store.verify_service_connector(
+            connector_details,
+            list_resources=True,
+        )
+        can_generate_long_tokens = not zen_store.get_service_connector(
+            connector_details
+        ).configuration.get("generate_temporary_tokens", True)
+    else:
+        resource_model = zen_store.verify_service_connector_config(
+            service_connector=ServiceConnectorRequest(
+                user=client.active_user.id,
+                workspace=client.active_workspace.id,
+                name="fake",
+                connector_type=connector_details.type,
+                auth_method=connector_details.auth_method,
+                configuration=connector_details.configuration,
+                secrets={},
+                labels={},
+            ),
+            list_resources=True,
+        )
+        can_generate_long_tokens = True
+
+    resources = resource_model.resources
+
+    if isinstance(
+        resource_model.connector_type,
+        str,
+    ):
+        connector_type = resource_model.connector_type
+    else:
+        connector_type = resource_model.connector_type.connector_type
+
+    artifact_stores: List[ResourcesInfo] = []
+    orchestrators: List[ResourcesInfo] = []
+    container_registries: List[ResourcesInfo] = []
+
+    if connector_type == "aws":
+        for each in resources:
+            if each.resource_ids:
+                if each.resource_type == "s3-bucket":
+                    artifact_stores.append(
+                        _prepare_resource_info(
+                            connector_details=connector_details,
+                            resource_ids=each.resource_ids,
+                            stack_component_type=StackComponentType.ARTIFACT_STORE,
+                            flavor="s3",
+                            required_configuration={"path": "Path"},
+                            use_resource_value_as_fixed_config=True,
+                            flavor_display_name="S3 Bucket",
+                        )
+                    )
+                if each.resource_type == "aws-generic":
+                    orchestrators.append(
+                        _prepare_resource_info(
+                            connector_details=connector_details,
+                            resource_ids=each.resource_ids,
+                            stack_component_type=StackComponentType.ORCHESTRATOR,
+                            flavor="sagemaker",
+                            required_configuration={
+                                "execution_role": "execution role ARN"
+                            },
+                            flavor_display_name="AWS Sagemaker",
+                        )
+                    )
+                    if can_generate_long_tokens:
+                        orchestrators.append(
+                            _prepare_resource_info(
+                                connector_details=connector_details,
+                                resource_ids=each.resource_ids,
+                                stack_component_type=StackComponentType.ORCHESTRATOR,
+                                flavor="vm_aws",
+                                required_configuration={"region": "region"},
+                                use_resource_value_as_fixed_config=True,
+                                flavor_display_name="Skypilot (EC2)",
+                            )
+                        )
+
+                if each.resource_type == "kubernetes-cluster":
+                    orchestrators.append(
+                        _prepare_resource_info(
+                            connector_details=connector_details,
+                            resource_ids=each.resource_ids,
+                            stack_component_type=StackComponentType.ORCHESTRATOR,
+                            flavor="kubernetes",
+                            required_configuration={},
+                            flavor_display_name="Kubernetes",
+                        )
+                    )
+                if each.resource_type == "docker-registry":
+                    container_registries.append(
+                        _prepare_resource_info(
+                            connector_details=connector_details,
+                            resource_ids=each.resource_ids,
+                            stack_component_type=StackComponentType.CONTAINER_REGISTRY,
+                            flavor="aws",
+                            required_configuration={"uri": "URI"},
+                            use_resource_value_as_fixed_config=True,
+                            flavor_display_name="ECR",
+                        )
+                    )
+
+    elif connector_type == "gcp":
+        for each in resources:
+            if each.resource_ids:
+                if each.resource_type == "gcs-bucket":
+                    artifact_stores.append(
+                        _prepare_resource_info(
+                            connector_details=connector_details,
+                            resource_ids=each.resource_ids,
+                            stack_component_type=StackComponentType.ARTIFACT_STORE,
+                            flavor="gcp",
+                            required_configuration={},
+                            flavor_display_name="GCS Bucket",
+                        )
+                    )
+                if each.resource_type == "gcp-generic":
+                    orchestrators.append(
+                        _prepare_resource_info(
+                            connector_details=connector_details,
+                            resource_ids=each.resource_ids,
+                            stack_component_type=StackComponentType.ORCHESTRATOR,
+                            flavor="vertex",
+                            required_configuration={"location": "region name"},
+                            flavor_display_name="Vertex AI",
+                        )
+                    )
+                    if can_generate_long_tokens:
+                        orchestrators.append(
+                            _prepare_resource_info(
+                                connector_details=connector_details,
+                                resource_ids=each.resource_ids,
+                                stack_component_type=StackComponentType.ORCHESTRATOR,
+                                flavor="vm_gcp",
+                                required_configuration={
+                                    "region": "region name"
+                                },
+                                flavor_display_name="Skypilot (Compute)",
+                            )
+                        )
+
+                if each.resource_type == "kubernetes-cluster":
+                    orchestrators.append(
+                        _prepare_resource_info(
+                            connector_details=connector_details,
+                            resource_ids=each.resource_ids,
+                            stack_component_type=StackComponentType.ORCHESTRATOR,
+                            flavor="kubernetes",
+                            required_configuration={},
+                            flavor_display_name="Kubernetes",
+                        )
+                    )
+                if each.resource_type == "docker-registry":
+                    container_registries.append(
+                        _prepare_resource_info(
+                            connector_details=connector_details,
+                            resource_ids=each.resource_ids,
+                            stack_component_type=StackComponentType.CONTAINER_REGISTRY,
+                            flavor="gcp",
+                            required_configuration={"uri": "URI"},
+                            use_resource_value_as_fixed_config=True,
+                            flavor_display_name="GCR",
+                        )
+                    )
+
+    elif connector_type == "azure":
+        pass
+
+    _raise_specific_cloud_exception_if_needed(
+        cloud_provider=connector_type,
+        artifact_stores=artifact_stores,
+        orchestrators=orchestrators,
+        container_registries=container_registries,
+    )
+
+    return ServiceConnectorResourcesInfo(
+        connector_type=connector_type,
+        components_resources_info={
+            StackComponentType.ARTIFACT_STORE: artifact_stores,
+            StackComponentType.ORCHESTRATOR: orchestrators,
+            StackComponentType.CONTAINER_REGISTRY: container_registries,
+        },
+    )