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

zenml/integrations/huggingface/materializers/huggingface_datasets_materializer.py

@@ -1,4 +1,4 @@
-#  Copyright (c) ZenML GmbH 2021. All Rights Reserved.
+#  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.
@@ -16,12 +16,21 @@
 import os
 from collections import defaultdict
 from tempfile import TemporaryDirectory, mkdtemp
-from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Dict,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+)
 
 from datasets import Dataset, load_from_disk
 from datasets.dataset_dict import DatasetDict
 
-from zenml.enums import ArtifactType
+from zenml.enums import ArtifactType, VisualizationType
 from zenml.io import fileio
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.materializers.pandas_materializer import PandasMaterializer
@@ -33,6 +42,31 @@ if TYPE_CHECKING:
 DEFAULT_DATASET_DIR = "hf_datasets"
 
 
+def extract_repo_name(checksum_str: str) -> Optional[str]:
+    """Extracts the repo name from the checksum string.
+
+    An example of a checksum_str is:
+    "hf://datasets/nyu-mll/glue@bcdcba79d07bc864c1c254ccfcedcce55bcc9a8c/mrpc/train-00000-of-00001.parquet"
+    and the expected output is "nyu-mll/glue".
+
+    Args:
+        checksum_str: The checksum_str to extract the repo name from.
+
+    Returns:
+        str: The extracted repo name.
+    """
+    dataset = None
+    try:
+        parts = checksum_str.split("/")
+        if len(parts) >= 4:
+            # Case: nyu-mll/glue
+            dataset = f"{parts[3]}/{parts[4].split('@')[0]}"
+    except Exception:  # pylint: disable=broad-except
+        pass
+
+    return dataset
+
+
 class HFDatasetMaterializer(BaseMaterializer):
     """Materializer to read data to and from huggingface datasets."""
 
@@ -103,3 +137,54 @@ class HFDatasetMaterializer(BaseMaterializer):
                     metadata[key][dataset_name] = value
             return dict(metadata)
         raise ValueError(f"Unsupported type {type(ds)}")
+
+    def save_visualizations(
+        self, ds: Union[Dataset, DatasetDict]
+    ) -> Dict[str, VisualizationType]:
+        """Save visualizations for the dataset.
+
+        Args:
+            ds: The Dataset or DatasetDict to visualize.
+
+        Returns:
+            A dictionary mapping visualization paths to their types.
+
+        Raises:
+            ValueError: If the given object is not a `Dataset` or `DatasetDict`.
+        """
+        visualizations = {}
+
+        if isinstance(ds, Dataset):
+            datasets = {"default": ds}
+        elif isinstance(ds, DatasetDict):
+            datasets = ds
+        else:
+            raise ValueError(f"Unsupported type {type(ds)}")
+
+        for name, dataset in datasets.items():
+            # Generate a unique identifier for the dataset
+            if dataset.info.download_checksums:
+                dataset_id = extract_repo_name(
+                    [x for x in dataset.info.download_checksums.keys()][0]
+                )
+                if dataset_id:
+                    # Create the iframe HTML
+                    html = f"""
+                    <iframe
+                    src="https://huggingface.co/datasets/{dataset_id}/embed/viewer"
+                    frameborder="0"
+                    width="100%"
+                    height="560px"
+                    ></iframe>
+                    """
+
+                    # Save the HTML to a file
+                    visualization_path = os.path.join(
+                        self.uri, f"{name}_viewer.html"
+                    )
+                    with fileio.open(visualization_path, "w") as f:
+                        f.write(html)
+
+                    visualizations[visualization_path] = VisualizationType.HTML
+
+        return visualizations

zenml/integrations/huggingface/steps/accelerate_runner.py

@@ -111,15 +111,9 @@ def run_with_accelerate(
                     if isinstance(v, bool):
                         if v:
                             commands.append(f"--{k}")
-                    elif isinstance(v, str):
-                        commands += [f"--{k}", '"{v}"']
                     elif type(v) in (list, tuple, set):
                         for each in v:
-                            commands.append(f"--{k}")
-                            if isinstance(each, str):
-                                commands.append(f'"{each}"')
-                            else:
-                                commands.append(f"{each}")
+                            commands += [f"--{k}", f"{each}"]
                     else:
                         commands += [f"--{k}", f"{v}"]
 

zenml/integrations/kubernetes/__init__.py

@@ -24,6 +24,7 @@ from zenml.integrations.integration import Integration
 from zenml.stack import Flavor
 
 KUBERNETES_ORCHESTRATOR_FLAVOR = "kubernetes"
+KUBERNETES_STEP_OPERATOR_FLAVOR = "kubernetes"
 
 
 class KubernetesIntegration(Integration):
@@ -42,10 +43,10 @@ class KubernetesIntegration(Integration):
             List of new stack component flavors.
         """
         from zenml.integrations.kubernetes.flavors import (
-            KubernetesOrchestratorFlavor,
+            KubernetesOrchestratorFlavor, KubernetesStepOperatorFlavor
         )
 
-        return [KubernetesOrchestratorFlavor]
+        return [KubernetesOrchestratorFlavor, KubernetesStepOperatorFlavor]
 
 
 KubernetesIntegration.check_installation()

zenml/integrations/kubernetes/flavors/__init__.py

@@ -18,9 +18,17 @@ from zenml.integrations.kubernetes.flavors.kubernetes_orchestrator_flavor import
     KubernetesOrchestratorFlavor,
     KubernetesOrchestratorSettings,
 )
+from zenml.integrations.kubernetes.flavors.kubernetes_step_operator_flavor import (
+    KubernetesStepOperatorConfig,
+    KubernetesStepOperatorFlavor,
+    KubernetesStepOperatorSettings,
+)
 
 __all__ = [
     "KubernetesOrchestratorFlavor",
     "KubernetesOrchestratorConfig",
     "KubernetesOrchestratorSettings",
+    "KubernetesStepOperatorConfig",
+    "KubernetesStepOperatorFlavor",
+    "KubernetesStepOperatorSettings",
 ]

zenml/integrations/kubernetes/flavors/kubernetes_step_operator_flavor.py

@@ -0,0 +1,166 @@
+#  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.
+"""Kubernetes step operator flavor."""
+
+from typing import TYPE_CHECKING, Optional, Type
+
+from zenml.config.base_settings import BaseSettings
+from zenml.constants import KUBERNETES_CLUSTER_RESOURCE_TYPE
+from zenml.integrations.kubernetes import KUBERNETES_STEP_OPERATOR_FLAVOR
+from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings
+from zenml.models import ServiceConnectorRequirements
+from zenml.step_operators import BaseStepOperatorConfig, BaseStepOperatorFlavor
+
+if TYPE_CHECKING:
+    from zenml.integrations.kubernetes.step_operators import (
+        KubernetesStepOperator,
+    )
+
+
+class KubernetesStepOperatorSettings(BaseSettings):
+    """Settings for the Kubernetes step operator.
+
+    Attributes:
+        pod_settings: Pod settings to apply to pods executing the steps.
+        service_account_name: Name of the service account to use for the pod.
+        privileged: If the container should be run in privileged mode.
+    """
+
+    pod_settings: Optional[KubernetesPodSettings] = None
+    service_account_name: Optional[str] = None
+    privileged: bool = False
+
+
+class KubernetesStepOperatorConfig(
+    BaseStepOperatorConfig, KubernetesStepOperatorSettings
+):
+    """Configuration for the Kubernetes step operator.
+
+    Attributes:
+        kubernetes_namespace: Name of the Kubernetes namespace to be used.
+        incluster: If `True`, the step operator will run the pipeline inside the
+            same cluster in which the orchestrator is running. For this to work,
+            the pod running the orchestrator needs permissions to create new
+            pods. If set, the `kubernetes_context` config option is ignored. If
+            the stack component is linked to a Kubernetes service connector,
+            this field is ignored.
+        kubernetes_context: Name of a Kubernetes context to run pipelines in.
+            If the stack component is linked to a Kubernetes service connector,
+            this field is ignored. Otherwise, it is mandatory.
+    """
+
+    kubernetes_namespace: str = "zenml"
+    incluster: bool = False
+    kubernetes_context: Optional[str] = None
+
+    @property
+    def is_remote(self) -> bool:
+        """Checks if this stack component is running remotely.
+
+        This designation is used to determine if the stack component can be
+        used with a local ZenML database or if it requires a remote ZenML
+        server.
+
+        Returns:
+            True if this config is for a remote component, False otherwise.
+        """
+        return True
+
+    @property
+    def is_local(self) -> bool:
+        """Checks if this stack component is running locally.
+
+        Returns:
+            True if this config is for a local component, False otherwise.
+        """
+        return False
+
+
+class KubernetesStepOperatorFlavor(BaseStepOperatorFlavor):
+    """Kubernetes step operator flavor."""
+
+    @property
+    def name(self) -> str:
+        """Name of the flavor.
+
+        Returns:
+            The name of the flavor.
+        """
+        return KUBERNETES_STEP_OPERATOR_FLAVOR
+
+    @property
+    def service_connector_requirements(
+        self,
+    ) -> Optional[ServiceConnectorRequirements]:
+        """Service connector resource requirements for service connectors.
+
+        Specifies resource requirements that are used to filter the available
+        service connector types that are compatible with this flavor.
+
+        Returns:
+            Requirements for compatible service connectors, if a service
+            connector is required for this flavor.
+        """
+        return ServiceConnectorRequirements(
+            resource_type=KUBERNETES_CLUSTER_RESOURCE_TYPE,
+        )
+
+    @property
+    def docs_url(self) -> Optional[str]:
+        """A url to point at docs explaining this flavor.
+
+        Returns:
+            A flavor docs url.
+        """
+        return self.generate_default_docs_url()
+
+    @property
+    def sdk_docs_url(self) -> Optional[str]:
+        """A url to point at SDK docs explaining this flavor.
+
+        Returns:
+            A flavor SDK docs url.
+        """
+        return self.generate_default_sdk_docs_url()
+
+    @property
+    def logo_url(self) -> str:
+        """A url to represent the flavor in the dashboard.
+
+        Returns:
+            The flavor logo.
+        """
+        return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/step_operator/kubernetes.png"
+
+    @property
+    def config_class(self) -> Type[KubernetesStepOperatorConfig]:
+        """Returns `KubernetesStepOperatorConfig` config class.
+
+        Returns:
+                The config class.
+        """
+        return KubernetesStepOperatorConfig
+
+    @property
+    def implementation_class(self) -> Type["KubernetesStepOperator"]:
+        """Implementation class for this flavor.
+
+        Returns:
+            The implementation class.
+        """
+        from zenml.integrations.kubernetes.step_operators import (
+            KubernetesStepOperator,
+        )
+
+        return KubernetesStepOperator

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/kubernetes/step_operators/__init__.py

@@ -0,0 +1,22 @@
+#  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.
+"""Kubernetes step operator."""
+
+from zenml.integrations.kubernetes.step_operators.kubernetes_step_operator  import (
+    KubernetesStepOperator,
+)
+
+__all__ = [
+    "KubernetesStepOperator",
+]

zenml/integrations/kubernetes/step_operators/kubernetes_step_operator.py

@@ -0,0 +1,235 @@
+#  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.
+"""Kubernetes step operator implementation."""
+
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type, cast
+
+from kubernetes import client as k8s_client
+
+from zenml.config.base_settings import BaseSettings
+from zenml.config.build_configuration import BuildConfiguration
+from zenml.enums import StackComponentType
+from zenml.integrations.kubernetes.flavors import (
+    KubernetesStepOperatorConfig,
+    KubernetesStepOperatorSettings,
+)
+from zenml.integrations.kubernetes.orchestrators import kube_utils
+from zenml.integrations.kubernetes.orchestrators.manifest_utils import (
+    build_pod_manifest,
+)
+from zenml.logger import get_logger
+from zenml.stack import Stack, StackValidator
+from zenml.step_operators import BaseStepOperator
+
+if TYPE_CHECKING:
+    from zenml.config.base_settings import BaseSettings
+    from zenml.config.step_run_info import StepRunInfo
+    from zenml.models import PipelineDeploymentBase
+
+logger = get_logger(__name__)
+
+KUBERNETES_STEP_OPERATOR_DOCKER_IMAGE_KEY = "kubernetes_step_operator"
+
+
+class KubernetesStepOperator(BaseStepOperator):
+    """Step operator to run on Kubernetes."""
+
+    _k8s_client: Optional[k8s_client.ApiClient] = None
+
+    @property
+    def config(self) -> KubernetesStepOperatorConfig:
+        """Returns the `KubernetesStepOperatorConfig` config.
+
+        Returns:
+            The configuration.
+        """
+        return cast(KubernetesStepOperatorConfig, self._config)
+
+    @property
+    def settings_class(self) -> Optional[Type["BaseSettings"]]:
+        """Settings class for the Kubernetes step operator.
+
+        Returns:
+            The settings class.
+        """
+        return KubernetesStepOperatorSettings
+
+    @property
+    def validator(self) -> Optional[StackValidator]:
+        """Validates the stack.
+
+        Returns:
+            A validator that checks that the stack contains a remote container
+            registry and a remote artifact store.
+        """
+
+        def _validate_remote_components(stack: "Stack") -> Tuple[bool, str]:
+            if stack.artifact_store.config.is_local:
+                return False, (
+                    "The Kubernetes step operator runs code remotely and "
+                    "needs to write files into the artifact store, but the "
+                    f"artifact store `{stack.artifact_store.name}` of the "
+                    "active stack is local. Please ensure that your stack "
+                    "contains a remote artifact store when using the Vertex "
+                    "step operator."
+                )
+
+            container_registry = stack.container_registry
+            assert container_registry is not None
+
+            if container_registry.config.is_local:
+                return False, (
+                    "The Kubernetes step operator runs code remotely and "
+                    "needs to push/pull Docker images, but the "
+                    f"container registry `{container_registry.name}` of the "
+                    "active stack is local. Please ensure that your stack "
+                    "contains a remote container registry when using the "
+                    "Kubernetes step operator."
+                )
+
+            return True, ""
+
+        return StackValidator(
+            required_components={
+                StackComponentType.CONTAINER_REGISTRY,
+                StackComponentType.IMAGE_BUILDER,
+            },
+            custom_validation_function=_validate_remote_components,
+        )
+
+    def get_docker_builds(
+        self, deployment: "PipelineDeploymentBase"
+    ) -> List["BuildConfiguration"]:
+        """Gets the Docker builds required for the component.
+
+        Args:
+            deployment: The pipeline deployment for which to get the builds.
+
+        Returns:
+            The required Docker builds.
+        """
+        builds = []
+        for step_name, step in deployment.step_configurations.items():
+            if step.config.step_operator == self.name:
+                build = BuildConfiguration(
+                    key=KUBERNETES_STEP_OPERATOR_DOCKER_IMAGE_KEY,
+                    settings=step.config.docker_settings,
+                    step_name=step_name,
+                )
+                builds.append(build)
+
+        return builds
+
+    def get_kube_client(self) -> k8s_client.ApiClient:
+        """Get the Kubernetes API client.
+
+        Returns:
+            The Kubernetes API client.
+
+        Raises:
+            RuntimeError: If the service connector returns an unexpected client.
+        """
+        if self.config.incluster:
+            kube_utils.load_kube_config(incluster=True)
+            self._k8s_client = k8s_client.ApiClient()
+            return self._k8s_client
+
+        # Refresh the client also if the connector has expired
+        if self._k8s_client and not self.connector_has_expired():
+            return self._k8s_client
+
+        connector = self.get_connector()
+        if connector:
+            client = connector.connect()
+            if not isinstance(client, k8s_client.ApiClient):
+                raise RuntimeError(
+                    f"Expected a k8s_client.ApiClient while trying to use the "
+                    f"linked connector, but got {type(client)}."
+                )
+            self._k8s_client = client
+        else:
+            kube_utils.load_kube_config(
+                context=self.config.kubernetes_context,
+            )
+            self._k8s_client = k8s_client.ApiClient()
+
+        return self._k8s_client
+
+    @property
+    def _k8s_core_api(self) -> k8s_client.CoreV1Api:
+        """Getter for the Kubernetes Core API client.
+
+        Returns:
+            The Kubernetes Core API client.
+        """
+        return k8s_client.CoreV1Api(self.get_kube_client())
+
+    def launch(
+        self,
+        info: "StepRunInfo",
+        entrypoint_command: List[str],
+        environment: Dict[str, str],
+    ) -> None:
+        """Launches a step on Kubernetes.
+
+        Args:
+            info: Information about the step run.
+            entrypoint_command: Command that executes the step.
+            environment: Environment variables to set in the step operator
+                environment.
+        """
+        settings = cast(
+            KubernetesStepOperatorSettings, self.get_settings(info)
+        )
+        image_name = info.get_image(
+            key=KUBERNETES_STEP_OPERATOR_DOCKER_IMAGE_KEY
+        )
+
+        pod_name = f"{info.run_name}_{info.pipeline_step_name}"
+        pod_name = kube_utils.sanitize_pod_name(pod_name)
+
+        command = entrypoint_command[:3]
+        args = entrypoint_command[3:]
+
+        # Create and run the orchestrator pod.
+        pod_manifest = build_pod_manifest(
+            run_name=info.run_name,
+            pod_name=pod_name,
+            pipeline_name=info.pipeline.name,
+            image_name=image_name,
+            command=command,
+            args=args,
+            privileged=settings.privileged,
+            service_account_name=settings.service_account_name,
+            pod_settings=settings.pod_settings,
+            env=environment,
+            mount_local_stores=False,
+        )
+
+        self._k8s_core_api.create_namespaced_pod(
+            namespace=self.config.kubernetes_namespace,
+            body=pod_manifest,
+        )
+
+        logger.info(
+            "Waiting for pod of step `%s` to start...", info.pipeline_step_name
+        )
+        kube_utils.wait_pod(
+            kube_client_fn=self.get_kube_client,
+            pod_name=pod_name,
+            namespace=self.config.kubernetes_namespace,
+            exit_condition_lambda=kube_utils.pod_is_done,
+            stream_logs=True,
+        )
+        logger.info("Pod of step `%s` completed.", info.pipeline_step_name)