zenml-nightly 0.70.0.dev20241122__py3-none-any.whl → 0.70.0.dev20241201__py3-none-any.whl

zenml/constants.py

@@ -353,6 +353,7 @@ FLAVORS = "/flavors"
 GET_OR_CREATE = "/get-or-create"
 HEALTH = "/health"
 INFO = "/info"
+LOAD_INFO = "/load-info"
 LOGIN = "/login"
 LOGOUT = "/logout"
 LOGS = "/logs"

zenml/enums.py

@@ -60,6 +60,7 @@ class VisualizationType(StrEnum):
     HTML = "html"
     IMAGE = "image"
     MARKDOWN = "markdown"
+    JSON = "json"
 
 
 class ZenMLServiceType(StrEnum):

zenml/image_builders/local_image_builder.py

@@ -68,15 +68,25 @@ class LocalImageBuilder(BaseImageBuilder):
             )
 
         if not docker_utils.check_docker():
+            # For 3., this is not supported by the python docker library
+            # https://github.com/docker/docker-py/issues/3146
             raise RuntimeError(
-                "Unable to connect to the Docker daemon. There are two "
+                "Unable to connect to the Docker daemon. There are three "
                 "common causes for this:\n"
                 "1) The Docker daemon isn't running.\n"
                 "2) The Docker client isn't configured correctly. The client "
                 "loads its configuration from the following file: "
                 "$HOME/.docker/config.json. If your configuration file is in a "
-                "different location, set it using the `DOCKER_CONFIG` "
-                "environment variable."
+                "different location, set the `DOCKER_CONFIG` environment "
+                "variable to the directory that contains your `config.json` "
+                "file.\n"
+                "3) If your Docker CLI is working fine but you ran into this "
+                "issue, you might be using a non-default Docker context which "
+                "is not supported by the Docker python library. To verify "
+                "this, run `docker context ls` and check which context has a "
+                "`*` next to it. If this is not the `default` context, copy "
+                "the `DOCKER ENDPOINT` value of that context and set the "
+                "`DOCKER_HOST` environment variable to that value."
             )
 
     def build(

zenml/integrations/__init__.py

@@ -48,6 +48,7 @@ from zenml.integrations.lightgbm import LightGBMIntegration  # noqa
 
 # from zenml.integrations.llama_index import LlamaIndexIntegration  # noqa
 from zenml.integrations.mlflow import MlflowIntegration  # noqa
+from zenml.integrations.modal import ModalIntegration  # noqa
 from zenml.integrations.neptune import NeptuneIntegration  # noqa
 from zenml.integrations.neural_prophet import NeuralProphetIntegration  # noqa
 from zenml.integrations.numpy import NumpyIntegration  # noqa

zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py

@@ -305,8 +305,21 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
             # Retrieve Executor arguments provided in the Step settings.
             if use_training_step:
                 args_for_step_executor = step_settings.estimator_args or {}
+                args_for_step_executor.setdefault(
+                    "volume_size", step_settings.volume_size_in_gb
+                )
+                args_for_step_executor.setdefault(
+                    "max_run", step_settings.max_runtime_in_seconds
+                )
             else:
                 args_for_step_executor = step_settings.processor_args or {}
+                args_for_step_executor.setdefault(
+                    "volume_size_in_gb", step_settings.volume_size_in_gb
+                )
+                args_for_step_executor.setdefault(
+                    "max_runtime_in_seconds",
+                    step_settings.max_runtime_in_seconds,
+                )
 
             # Set default values from configured orchestrator Component to
             # arguments to be used when they are not present in processor_args.
@@ -314,12 +327,7 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
                 "role",
                 step_settings.execution_role or self.config.execution_role,
             )
-            args_for_step_executor.setdefault(
-                "volume_size_in_gb", step_settings.volume_size_in_gb
-            )
-            args_for_step_executor.setdefault(
-                "max_runtime_in_seconds", step_settings.max_runtime_in_seconds
-            )
+
             tags = step_settings.tags
             args_for_step_executor.setdefault(
                 "tags",

zenml/integrations/constants.py

@@ -42,6 +42,7 @@ LANGCHAIN = "langchain"
 LIGHTGBM = "lightgbm"
 # LLAMA_INDEX = "llama_index"
 MLFLOW = "mlflow"
+MODAL = "modal"
 NEPTUNE = "neptune"
 NEURAL_PROPHET = "neural_prophet"
 NUMPY = "numpy"

zenml/integrations/feast/__init__.py

@@ -31,7 +31,7 @@ class FeastIntegration(Integration):
 
     NAME = FEAST
     # click is added to keep the feast click version in sync with ZenML's click
-    REQUIREMENTS = ["feast", "click>=8.0.1,<8.1.4"]
+    REQUIREMENTS = ["feast>=0.12.0", "click>=8.0.1,<8.1.4"]
     REQUIREMENTS_IGNORED_ON_UNINSTALL = ["click", "pandas"]
 
     @classmethod

zenml/integrations/feast/feature_stores/feast_feature_store.py

@@ -16,7 +16,7 @@
 from typing import Any, Dict, List, Union, cast
 
 import pandas as pd
-from feast import FeatureStore  # type: ignore
+from feast import FeatureService, FeatureStore  # type: ignore
 from feast.infra.registry.base_registry import BaseRegistry  # type: ignore
 
 from zenml.feature_stores.base_feature_store import BaseFeatureStore
@@ -43,14 +43,14 @@ class FeastFeatureStore(BaseFeatureStore):
     def get_historical_features(
         self,
         entity_df: Union[pd.DataFrame, str],
-        features: List[str],
+        features: Union[List[str], FeatureService],
         full_feature_names: bool = False,
     ) -> pd.DataFrame:
         """Returns the historical features for training or batch scoring.
 
         Args:
             entity_df: The entity DataFrame or entity name.
-            features: The features to retrieve.
+            features: The features to retrieve or a FeatureService.
             full_feature_names: Whether to return the full feature names.
 
         Raise:
@@ -70,14 +70,14 @@ class FeastFeatureStore(BaseFeatureStore):
     def get_online_features(
         self,
         entity_rows: List[Dict[str, Any]],
-        features: List[str],
+        features: Union[List[str], FeatureService],
         full_feature_names: bool = False,
     ) -> Dict[str, Any]:
         """Returns the latest online feature data.
 
         Args:
             entity_rows: The entity rows to retrieve.
-            features: The features to retrieve.
+            features: The features to retrieve or a FeatureService.
             full_feature_names: Whether to return the full feature names.
 
         Raise:
@@ -118,17 +118,21 @@ class FeastFeatureStore(BaseFeatureStore):
         fs = FeatureStore(repo_path=self.config.feast_repo)
         return [ds.name for ds in fs.list_entities()]
 
-    def get_feature_services(self) -> List[str]:
-        """Returns the feature service names.
+    def get_feature_services(self) -> List[FeatureService]:
+        """Returns the feature services.
 
         Raise:
             ConnectionError: If the online component (Redis) is not available.
 
         Returns:
-            The feature service names.
+            The feature services.
         """
         fs = FeatureStore(repo_path=self.config.feast_repo)
-        return [ds.name for ds in fs.list_feature_services()]
+        feature_services: List[FeatureService] = list(
+            fs.list_feature_services()
+        )
+
+        return feature_services
 
     def get_feature_views(self) -> List[str]:
         """Returns the feature view names.

zenml/integrations/kubernetes/orchestrators/kube_utils.py

@@ -42,8 +42,8 @@ from kubernetes import config as k8s_config
 from kubernetes.client.rest import ApiException
 
 from zenml.integrations.kubernetes.orchestrators.manifest_utils import (
-    build_cluster_role_binding_manifest_for_service_account,
     build_namespace_manifest,
+    build_role_binding_manifest_for_service_account,
     build_service_account_manifest,
 )
 from zenml.logger import get_logger
@@ -94,18 +94,62 @@ def load_kube_config(
         k8s_config.load_kube_config(context=context)
 
 
-def sanitize_pod_name(pod_name: str) -> str:
+def calculate_max_pod_name_length_for_namespace(namespace: str) -> int:
+    """Calculate the max pod length for a certain namespace.
+
+    Args:
+        namespace: The namespace in which the pod will be created.
+
+    Returns:
+        The maximum pod name length.
+    """
+    # Kubernetes allows Pod names to have 253 characters. However, when
+    # creating a pod they try to create a log file which is called
+    # <NAMESPACE>_<POD_NAME>_<UUID>, which adds additional characters and
+    # runs into filesystem limitations for filename lengths (255). We therefore
+    # subtract the length of a UUID (36), the two underscores and the
+    # namespace length from the max filename length.
+    return 255 - 38 - len(namespace)
+
+
+def sanitize_pod_name(pod_name: str, namespace: str) -> str:
     """Sanitize pod names so they conform to Kubernetes pod naming convention.
 
     Args:
         pod_name: Arbitrary input pod name.
+        namespace: Namespace in which the Pod will be created.
 
     Returns:
         Sanitized pod name.
     """
+    # https://kubernetes.io/docs/concepts/overview/working-with-objects/names/#dns-subdomain-names
     pod_name = re.sub(r"[^a-z0-9-]", "-", pod_name.lower())
     pod_name = re.sub(r"^[-]+", "", pod_name)
-    return re.sub(r"[-]+", "-", pod_name)
+    pod_name = re.sub(r"[-]+$", "", pod_name)
+    pod_name = re.sub(r"[-]+", "-", pod_name)
+
+    allowed_length = calculate_max_pod_name_length_for_namespace(
+        namespace=namespace
+    )
+    return pod_name[:allowed_length]
+
+
+def sanitize_label(label: str) -> str:
+    """Sanitize a label for a Kubernetes resource.
+
+    Args:
+        label: The label to sanitize.
+
+    Returns:
+        The sanitized label.
+    """
+    # https://kubernetes.io/docs/concepts/overview/working-with-objects/names/#rfc-1035-label-names
+    label = re.sub(r"[^a-z0-9-]", "-", label.lower())
+    label = re.sub(r"^[-]+", "", label)
+    label = re.sub(r"[-]+$", "", label)
+    label = re.sub(r"[-]+", "-", label)
+
+    return label[:63]
 
 
 def pod_is_not_pending(pod: k8s_client.V1Pod) -> bool:
@@ -288,7 +332,7 @@ def create_edit_service_account(
     rbac_api: k8s_client.RbacAuthorizationV1Api,
     service_account_name: str,
     namespace: str,
-    cluster_role_binding_name: str = "zenml-edit",
+    role_binding_name: str = "zenml-edit",
 ) -> None:
     """Create a new Kubernetes service account with "edit" rights.
 
@@ -297,16 +341,17 @@ def create_edit_service_account(
         rbac_api: Client of Rbac Authorization V1 API of Kubernetes API.
         service_account_name: Name of the service account.
         namespace: Kubernetes namespace. Defaults to "default".
-        cluster_role_binding_name: Name of the cluster role binding.
-            Defaults to "zenml-edit".
+        role_binding_name: Name of the role binding. Defaults to "zenml-edit".
     """
-    crb_manifest = build_cluster_role_binding_manifest_for_service_account(
-        name=cluster_role_binding_name,
+    rb_manifest = build_role_binding_manifest_for_service_account(
+        name=role_binding_name,
         role_name="edit",
         service_account_name=service_account_name,
         namespace=namespace,
     )
-    _if_not_exists(rbac_api.create_cluster_role_binding)(body=crb_manifest)
+    _if_not_exists(rbac_api.create_namespaced_role_binding)(
+        namespace=namespace, body=rb_manifest
+    )
 
     sa_manifest = build_service_account_manifest(
         name=service_account_name, namespace=namespace

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -59,6 +59,7 @@ from zenml.integrations.kubernetes.orchestrators.manifest_utils import (
     build_cron_job_manifest,
     build_pod_manifest,
 )
+from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings
 from zenml.logger import get_logger
 from zenml.orchestrators import ContainerizedOrchestrator
 from zenml.orchestrators.utils import get_orchestrator_run_name
@@ -328,6 +329,45 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
             custom_validation_function=_validate_local_requirements,
         )
 
+    @classmethod
+    def apply_default_resource_requests(
+        cls,
+        memory: str,
+        cpu: Optional[str] = None,
+        pod_settings: Optional[KubernetesPodSettings] = None,
+    ) -> KubernetesPodSettings:
+        """Applies default resource requests to a pod settings object.
+
+        Args:
+            memory: The memory resource request.
+            cpu: The CPU resource request.
+            pod_settings: The pod settings to update. A new one will be created
+                if not provided.
+
+        Returns:
+            The new or updated pod settings.
+        """
+        resources = {
+            "requests": {"memory": memory},
+        }
+        if cpu:
+            resources["requests"]["cpu"] = cpu
+        if not pod_settings:
+            pod_settings = KubernetesPodSettings(resources=resources)
+        elif not pod_settings.resources:
+            # We can't update the pod settings in place (because it's a frozen
+            # pydantic model), so we have to create a new one.
+            pod_settings = KubernetesPodSettings(
+                **pod_settings.model_dump(exclude_unset=True),
+                resources=resources,
+            )
+        else:
+            set_requests = pod_settings.resources.get("requests", {})
+            resources["requests"].update(set_requests)
+            pod_settings.resources["requests"] = resources["requests"]
+
+        return pod_settings
+
     def prepare_or_run_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
@@ -355,8 +395,19 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                 )
 
         pipeline_name = deployment.pipeline_configuration.name
-        orchestrator_run_name = get_orchestrator_run_name(pipeline_name)
-        pod_name = kube_utils.sanitize_pod_name(orchestrator_run_name)
+
+        # We already make sure the orchestrator run name has the correct length
+        # to make sure we don't cut off the randomized suffix later when
+        # sanitizing the pod name. This avoids any pod naming collisions.
+        max_length = kube_utils.calculate_max_pod_name_length_for_namespace(
+            namespace=self.config.kubernetes_namespace
+        )
+        orchestrator_run_name = get_orchestrator_run_name(
+            pipeline_name, max_length=max_length
+        )
+        pod_name = kube_utils.sanitize_pod_name(
+            orchestrator_run_name, namespace=self.config.kubernetes_namespace
+        )
 
         assert stack.container_registry
 
@@ -422,6 +473,17 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
             )
             return
 
+        # We set some default minimum resource requests for the orchestrator pod
+        # here if the user has not specified any, because the orchestrator pod
+        # takes up some memory resources itself and, if not specified, the pod
+        # will be scheduled on any node regardless of available memory and risk
+        # negatively impacting or even crashing the node due to memory pressure.
+        orchestrator_pod_settings = self.apply_default_resource_requests(
+            memory="400Mi",
+            cpu="100m",
+            pod_settings=settings.orchestrator_pod_settings,
+        )
+
         # Create and run the orchestrator pod.
         pod_manifest = build_pod_manifest(
             run_name=orchestrator_run_name,
@@ -431,7 +493,7 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
             command=command,
             args=args,
             privileged=False,
-            pod_settings=settings.orchestrator_pod_settings,
+            pod_settings=orchestrator_pod_settings,
             service_account_name=service_account_name,
             env=environment,
             mount_local_stores=self.config.is_local,

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator_entrypoint.py

@@ -90,7 +90,9 @@ def main() -> None:
         """
         # Define Kubernetes pod name.
         pod_name = f"{orchestrator_run_id}-{step_name}"
-        pod_name = kube_utils.sanitize_pod_name(pod_name)
+        pod_name = kube_utils.sanitize_pod_name(
+            pod_name, namespace=args.kubernetes_namespace
+        )
 
         image = KubernetesOrchestrator.get_image(
             deployment=deployment_config, step_name=step_name
@@ -116,6 +118,16 @@ def main() -> None:
         env = get_config_environment_vars()
         env[ENV_ZENML_KUBERNETES_RUN_ID] = orchestrator_run_id
 
+        # We set some default minimum memory resource requests for the step pod
+        # here if the user has not specified any, because the step pod takes up
+        # some memory resources itself and, if not specified, the pod will be
+        # scheduled on any node regardless of available memory and risk
+        # negatively impacting or even crashing the node due to memory pressure.
+        pod_settings = KubernetesOrchestrator.apply_default_resource_requests(
+            memory="400Mi",
+            pod_settings=settings.pod_settings,
+        )
+
         # Define Kubernetes pod manifest.
         pod_manifest = build_pod_manifest(
             pod_name=pod_name,
@@ -126,7 +138,7 @@ def main() -> None:
             args=step_args,
             env=env,
             privileged=settings.privileged,
-            pod_settings=settings.pod_settings,
+            pod_settings=pod_settings,
             service_account_name=settings.step_pod_service_account_name
             or settings.service_account_name,
             mount_local_stores=mount_local_stores,

zenml/integrations/kubernetes/orchestrators/manifest_utils.py

@@ -25,6 +25,7 @@ from zenml.constants import ENV_ZENML_ENABLE_REPO_INIT_WARNINGS
 from zenml.integrations.airflow.orchestrators.dag_generator import (
     ENV_ZENML_LOCAL_STORES_PATH,
 )
+from zenml.integrations.kubernetes.orchestrators import kube_utils
 from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings
 
 
@@ -167,8 +168,8 @@ def build_pod_manifest(
     # Add run_name and pipeline_name to the labels
     labels.update(
         {
-            "run": run_name,
-            "pipeline": pipeline_name,
+            "run": kube_utils.sanitize_label(run_name),
+            "pipeline": kube_utils.sanitize_label(pipeline_name),
         }
     )
 
@@ -304,13 +305,13 @@ def build_cron_job_manifest(
     return job_manifest
 
 
-def build_cluster_role_binding_manifest_for_service_account(
+def build_role_binding_manifest_for_service_account(
     name: str,
     role_name: str,
     service_account_name: str,
     namespace: str = "default",
 ) -> Dict[str, Any]:
-    """Build a manifest for a cluster role binding of a service account.
+    """Build a manifest for a role binding of a service account.
 
     Args:
         name: Name of the cluster role binding.
@@ -323,7 +324,7 @@ def build_cluster_role_binding_manifest_for_service_account(
     """
     return {
         "apiVersion": "rbac.authorization.k8s.io/v1",
-        "kind": "ClusterRoleBinding",
+        "kind": "RoleBinding",
         "metadata": {"name": name},
         "subjects": [
             {

zenml/integrations/kubernetes/service_connectors/kubernetes_service_connector.py

@@ -26,6 +26,7 @@ from typing import Any, List, Optional
 from kubernetes import client as k8s_client
 from kubernetes import config as k8s_config
 from pydantic import Field
+from urllib3.exceptions import HTTPError
 
 from zenml.constants import KUBERNETES_CLUSTER_RESOURCE_TYPE
 from zenml.exceptions import AuthorizationException
@@ -572,7 +573,7 @@ class KubernetesServiceConnector(ServiceConnector):
                 auth_settings=["BearerToken"],
                 response_type="VersionInfo",
             )
-        except k8s_client.ApiException as err:
+        except (k8s_client.ApiException, HTTPError) as err:
             raise AuthorizationException(
                 f"failed to verify Kubernetes cluster access: {err}"
             ) from err

zenml/integrations/kubernetes/step_operators/kubernetes_step_operator.py

@@ -197,7 +197,9 @@ class KubernetesStepOperator(BaseStepOperator):
         )
 
         pod_name = f"{info.run_name}_{info.pipeline_step_name}"
-        pod_name = kube_utils.sanitize_pod_name(pod_name)
+        pod_name = kube_utils.sanitize_pod_name(
+            pod_name, namespace=self.config.kubernetes_namespace
+        )
 
         command = entrypoint_command[:3]
         args = entrypoint_command[3:]

zenml/integrations/modal/__init__.py

@@ -0,0 +1,46 @@
+#  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.
+"""Modal integration for cloud-native step execution.
+
+The Modal integration sub-module provides a step operator flavor that allows
+executing steps on Modal's cloud infrastructure.
+"""
+from typing import List, Type
+
+from zenml.integrations.constants import MODAL
+from zenml.integrations.integration import Integration
+from zenml.stack import Flavor
+
+MODAL_STEP_OPERATOR_FLAVOR = "modal"
+
+
+class ModalIntegration(Integration):
+    """Definition of Modal integration for ZenML."""
+
+    NAME = MODAL
+    REQUIREMENTS = ["modal>=0.64.49,<1"]
+
+    @classmethod
+    def flavors(cls) -> List[Type[Flavor]]:
+        """Declare the stack component flavors for the Modal integration.
+
+        Returns:
+            List of new stack component flavors.
+        """
+        from zenml.integrations.modal.flavors import ModalStepOperatorFlavor
+
+        return [ModalStepOperatorFlavor]
+
+
+ModalIntegration.check_installation()

zenml/integrations/modal/flavors/__init__.py

@@ -0,0 +1,26 @@
+#  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.
+"""Modal integration flavors."""
+
+from zenml.integrations.modal.flavors.modal_step_operator_flavor import (
+    ModalStepOperatorConfig,
+    ModalStepOperatorFlavor,
+    ModalStepOperatorSettings,
+)
+
+__all__ = [
+    "ModalStepOperatorConfig",
+    "ModalStepOperatorFlavor",
+    "ModalStepOperatorSettings",
+]