zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/integrations/kubernetes/step_operators/kubernetes_step_operator.py

@@ -0,0 +1,236 @@
+#  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.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, namespace=self.config.kubernetes_namespace
+        )
+
+        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)

zenml/integrations/label_studio/__init__.py

@@ -14,7 +14,6 @@
 """Initialization of the Label Studio integration."""
 from typing import List, Type
 
-from zenml.enums import StackComponentType
 from zenml.integrations.constants import LABEL_STUDIO
 from zenml.integrations.integration import Integration
 from zenml.stack import Flavor
@@ -27,8 +26,7 @@ class LabelStudioIntegration(Integration):
 
     NAME = LABEL_STUDIO
     REQUIREMENTS = [
-        "label-studio>=1.6.0,<=1.9.1",
-        "label-studio-sdk>=0.0.17,<=0.0.32",
+        "label-studio-sdk>=1.0.0",
     ]
 
     @classmethod

zenml/integrations/label_studio/annotators/label_studio_annotator.py

@@ -17,13 +17,17 @@ import json
 import os
 import webbrowser
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple, cast
+from typing import Any, Dict, List, Optional, Tuple, Type, cast
 
-from label_studio_sdk import Client, Project  # type: ignore[import-not-found]
+from label_studio_sdk import Client, Project
 
+from zenml import get_step_context
 from zenml.annotators.base_annotator import BaseAnnotator
 from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
 from zenml.config.global_config import GlobalConfiguration
+from zenml.integrations.label_studio.flavors import (
+    LabelStudioAnnotatorSettings,
+)
 from zenml.integrations.label_studio.flavors.label_studio_annotator_flavor import (
     LabelStudioAnnotatorConfig,
 )
@@ -49,13 +53,26 @@ class LabelStudioAnnotator(BaseAnnotator, AuthenticationMixin):
         """
         return cast(LabelStudioAnnotatorConfig, self._config)
 
+    @property
+    def settings_class(self) -> Type[LabelStudioAnnotatorSettings]:
+        """Settings class for the Label Studio annotator.
+
+        Returns:
+            The settings class.
+        """
+        return LabelStudioAnnotatorSettings
+
     def get_url(self) -> str:
         """Gets the top-level URL of the annotation interface.
 
         Returns:
             The URL of the annotation interface.
         """
-        return f"{self.config.instance_url}:{self.config.port}"
+        return (
+            f"{self.config.instance_url}:{self.config.port}"
+            if self.config.port
+            else self.config.instance_url
+        )
 
     def get_url_for_dataset(self, dataset_name: str) -> str:
         """Gets the URL of the annotation interface for the given dataset.
@@ -126,14 +143,14 @@ class LabelStudioAnnotator(BaseAnnotator, AuthenticationMixin):
             f"`zenml annotator dataset list` to list all available datasets."
         )
 
-    def launch(self, url: Optional[str]) -> None:
+    def launch(self, **kwargs: Any) -> None:
         """Launches the annotation interface.
 
         Args:
-            url: The URL of the annotation interface.
+            **kwargs: Additional keyword arguments to pass to the
+                annotation client.
         """
-        if not url:
-            url = self.get_url()
+        url = kwargs.get("url") or self.get_url()
         if self._connection_available():
             webbrowser.open(url, new=1, autoraise=True)
         else:
@@ -151,12 +168,22 @@ class LabelStudioAnnotator(BaseAnnotator, AuthenticationMixin):
         Raises:
             ValueError: when unable to access the Label Studio API key.
         """
-        secret = self.get_authentication_secret()
-        if not secret:
-            raise ValueError(
-                "Unable to access predefined secret to access Label Studio API key."
+        try:
+            settings = cast(
+                LabelStudioAnnotatorSettings,
+                self.get_settings(get_step_context().step_run),
             )
-        api_key = secret.secret_values.get("api_key")
+            if settings.api_key is None:
+                raise RuntimeError
+            else:
+                api_key = settings.api_key
+        except RuntimeError:
+            if secret := self.get_authentication_secret():
+                api_key = secret.secret_values.get("api_key", "")
+            else:
+                raise ValueError(
+                    "Unable to access predefined secret to access Label Studio API key."
+                )
         if not api_key:
             raise ValueError(
                 "Unable to access Label Studio API key from secret."
@@ -500,6 +527,7 @@ class LabelStudioAnnotator(BaseAnnotator, AuthenticationMixin):
                 aws_access_key_id,
                 aws_secret_access_key,
                 aws_session_token,
+                _,
             ) = artifact_store.get_credentials()
 
             if aws_access_key_id and aws_secret_access_key:
@@ -549,9 +577,13 @@ class LabelStudioAnnotator(BaseAnnotator, AuthenticationMixin):
                 file_path = Path(
                     secret_folder, "google_application_credentials.json"
                 )
-                with open(file_path, "w") as f:
+                with os.fdopen(
+                    os.open(
+                        file_path, flags=os.O_RDWR | os.O_CREAT, mode=0o600
+                    ),
+                    "w",
+                ) as f:
                     f.write(json.dumps(gcp_credentials))
-                    file_path.chmod(0o600)
 
                 params.google_application_credentials = str(file_path)
 

zenml/integrations/label_studio/flavors/__init__.py

@@ -16,9 +16,11 @@
 from zenml.integrations.label_studio.flavors.label_studio_annotator_flavor import (
     LabelStudioAnnotatorConfig,
     LabelStudioAnnotatorFlavor,
+    LabelStudioAnnotatorSettings,
 )
 
 __all__ = [
     "LabelStudioAnnotatorFlavor",
     "LabelStudioAnnotatorConfig",
+    "LabelStudioAnnotatorSettings"
 ]

zenml/integrations/label_studio/flavors/label_studio_annotator_flavor.py

@@ -19,8 +19,10 @@ from zenml.annotators.base_annotator import (
     BaseAnnotatorConfig,
     BaseAnnotatorFlavor,
 )
+from zenml.config.base_settings import BaseSettings
 from zenml.integrations.label_studio import LABEL_STUDIO_ANNOTATOR_FLAVOR
 from zenml.stack.authentication_mixin import AuthenticationConfigMixin
+from zenml.utils.secret_utils import SecretField
 
 if TYPE_CHECKING:
     from zenml.integrations.label_studio.annotators import LabelStudioAnnotator
@@ -30,18 +32,31 @@ DEFAULT_LOCAL_INSTANCE_URL = "http://localhost"
 DEFAULT_LOCAL_LABEL_STUDIO_PORT = 8093
 
 
-class LabelStudioAnnotatorConfig(
-    BaseAnnotatorConfig, AuthenticationConfigMixin
-):
-    """Config for the Label Studio annotator.
+class LabelStudioAnnotatorSettings(BaseSettings):
+    """Label studio annotator settings.
 
     Attributes:
         instance_url: URL of the Label Studio instance.
         port: The port to use for the annotation interface.
+        api_key: The api_key for label studio.
     """
 
     instance_url: str = DEFAULT_LOCAL_INSTANCE_URL
-    port: int = DEFAULT_LOCAL_LABEL_STUDIO_PORT
+    port: Optional[int] = DEFAULT_LOCAL_LABEL_STUDIO_PORT
+    api_key: Optional[str] = SecretField(default=None)
+
+
+class LabelStudioAnnotatorConfig(
+    BaseAnnotatorConfig,
+    LabelStudioAnnotatorSettings,
+    AuthenticationConfigMixin,
+):
+    """Config for the Label Studio annotator.
+
+    This class combines settings and authentication configurations for
+    Label Studio into a single, usable configuration object without adding
+    additional functionality.
+    """
 
 
 class LabelStudioAnnotatorFlavor(BaseAnnotatorFlavor):

zenml/integrations/langchain/__init__.py

@@ -25,7 +25,13 @@ class LangchainIntegration(Integration):
     """Definition of langchain integration for ZenML."""
 
     NAME = LANGCHAIN
-    REQUIREMENTS = ["langchain==0.0.325", "pyyaml>=6.0.1"]
+    REQUIREMENTS = [
+        "langchain==0.3.0",
+        "langchain-community",
+        "pyyaml>=6.0.1",
+        "tenacity!=8.4.0",  # https://github.com/jd/tenacity/issues/471
+    ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["pyyaml","tenacity"]
 
     @classmethod
     def activate(cls) -> None:

zenml/integrations/langchain/materializers/document_materializer.py

@@ -13,20 +13,56 @@
 #  permissions and limitations under the License.
 """Implementation of ZenML's Langchain Document materializer."""
 
-import sys
-from typing import TYPE_CHECKING, Any, ClassVar, Tuple, Type
+import os
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type
+
+from langchain.docstore.document import Document
 
 from zenml.enums import ArtifactType
-from zenml.materializers.pydantic_materializer import PydanticMaterializer
+from zenml.materializers.base_materializer import BaseMaterializer
+from zenml.utils import yaml_utils
+
+if TYPE_CHECKING:
+    from zenml.metadata.metadata_types import MetadataType
 
-if TYPE_CHECKING and sys.version_info < (3, 8):
-    Document = Any
-else:
-    from langchain.docstore.document import Document
+DEFAULT_FILENAME = "data.json"
 
 
-class LangchainDocumentMaterializer(PydanticMaterializer):
+class LangchainDocumentMaterializer(BaseMaterializer):
     """Handle Langchain Document objects."""
 
     ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.DATA
     ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (Document,)
+
+    def load(self, data_type: Type["Document"]) -> Any:
+        """Reads BaseModel from JSON.
+
+        Args:
+            data_type: The type of the data to read.
+
+        Returns:
+            The data read.
+        """
+        data_path = os.path.join(self.uri, DEFAULT_FILENAME)
+        contents = yaml_utils.read_json(data_path)
+        return data_type.parse_raw(contents)
+
+    def save(self, data: "Document") -> None:
+        """Serialize a BaseModel to JSON.
+
+        Args:
+            data: The data to store.
+        """
+        data_path = os.path.join(self.uri, DEFAULT_FILENAME)
+        yaml_utils.write_json(data_path, data.json())
+
+    def extract_metadata(self, data: Document) -> Dict[str, "MetadataType"]:
+        """Extract metadata from the given BaseModel object.
+
+        Args:
+            data: The BaseModel object to extract metadata from.
+
+        Returns:
+            The extracted metadata as a dictionary.
+        """
+        return {"schema": data.schema()}

zenml/integrations/langchain/materializers/openai_embedding_materializer.py

@@ -24,11 +24,37 @@ from zenml.materializers.cloudpickle_materializer import (
 if TYPE_CHECKING and sys.version_info < (3, 8):
     OpenAIEmbeddings = Any
 else:
-    from langchain.embeddings import OpenAIEmbeddings
+    from langchain_community.embeddings import (
+        OpenAIEmbeddings,
+    )
 
 
 class LangchainOpenaiEmbeddingMaterializer(CloudpickleMaterializer):
-    """Handle langchain OpenAI embedding objects."""
+    """Materializer for Langchain OpenAI Embeddings."""
 
     ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.MODEL
     ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (OpenAIEmbeddings,)
+
+    def save(self, embeddings: Any) -> None:
+        """Saves the embeddings model after clearing non-picklable clients.
+
+        Args:
+            embeddings: The embeddings model to save.
+        """
+        # Clear the clients which will be recreated on load
+        embeddings.client = None
+        embeddings.async_client = None
+
+        # Use the parent class's save implementation which uses cloudpickle
+        super().save(embeddings)
+
+    def load(self, data_type: Type[Any]) -> Any:
+        """Loads the embeddings model and lets it recreate clients when needed.
+
+        Args:
+            data_type: The type of the data to load.
+
+        Returns:
+            The loaded embeddings model.
+        """
+        return super().load(data_type)

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/lightgbm/materializers/lightgbm_booster_materializer.py

@@ -14,7 +14,6 @@
 """Implementation of the LightGBM booster materializer."""
 
 import os
-import tempfile
 from typing import Any, ClassVar, Tuple, Type
 
 import lightgbm as lgb
@@ -42,18 +41,13 @@ class LightGBMBoosterMaterializer(BaseMaterializer):
             A lightgbm Booster object.
         """
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
 
-        # Create a temporary folder
-        temp_dir = tempfile.mkdtemp(prefix="zenml-temp-")
-        temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
-
-        # Copy from artifact store to temporary file
-        fileio.copy(filepath, temp_file)
-        booster = lgb.Booster(model_file=temp_file)
-
-        # Cleanup and return
-        fileio.rmtree(temp_dir)
-        return booster
+            # Copy from artifact store to temporary file
+            fileio.copy(filepath, temp_file)
+            booster = lgb.Booster(model_file=temp_file)
+            return booster
 
     def save(self, booster: lgb.Booster) -> None:
         """Creates a JSON serialization for a lightgbm Booster model.
@@ -62,8 +56,7 @@ class LightGBMBoosterMaterializer(BaseMaterializer):
             booster: A lightgbm Booster model.
         """
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
-
-        with tempfile.TemporaryDirectory() as tmp_dir:
-            tmp_path = os.path.join(tmp_dir, "model.txt")
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            tmp_path = os.path.join(temp_dir, "model.txt")
             booster.save_model(tmp_path)
             fileio.copy(tmp_path, filepath)

zenml/integrations/lightgbm/materializers/lightgbm_dataset_materializer.py

@@ -14,7 +14,6 @@
 """Implementation of the LightGBM materializer."""
 
 import os
-import tempfile
 from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type
 
 import lightgbm as lgb
@@ -46,16 +45,14 @@ class LightGBMDatasetMaterializer(BaseMaterializer):
         """
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
 
-        # Create a temporary folder
-        temp_dir = tempfile.mkdtemp(prefix="zenml-temp-")
-        temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
+        with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
+            temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
 
-        # Copy from artifact store to temporary file
-        fileio.copy(filepath, temp_file)
-        matrix = lgb.Dataset(temp_file, free_raw_data=False)
+            # Copy from artifact store to temporary file
+            fileio.copy(filepath, temp_file)
+            matrix = lgb.Dataset(temp_file, free_raw_data=False)
 
-        # No clean up this time because matrix is lazy loaded
-        return matrix
+            return matrix
 
     def save(self, matrix: lgb.Dataset) -> None:
         """Creates a binary serialization for a lightgbm.Dataset object.
@@ -65,14 +62,12 @@ class LightGBMDatasetMaterializer(BaseMaterializer):
         """
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
 
-        # Make a temporary phantom artifact
-        temp_dir = tempfile.mkdtemp(prefix="zenml-temp-")
-        temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
-        matrix.save_binary(temp_file)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
+            matrix.save_binary(temp_file)
 
-        # Copy it into artifact store
-        fileio.copy(temp_file, filepath)
-        fileio.rmtree(temp_dir)
+            # Copy it into artifact store
+            fileio.copy(temp_file, filepath)
 
     def extract_metadata(
         self, matrix: lgb.Dataset