zenml-nightly 0.70.0.dev20241126__py3-none-any.whl → 0.70.0.dev20241128__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.70.0.dev20241126
+0.70.0.dev20241128

zenml/artifact_stores/base_artifact_store.py

@@ -267,11 +267,11 @@ class BaseArtifactStore(StackComponent):
 
     # --- User interface ---
     @abstractmethod
-    def open(self, name: PathType, mode: str = "r") -> Any:
+    def open(self, path: PathType, mode: str = "r") -> Any:
         """Open a file at the given path.
 
         Args:
-            name: The path of the file to open.
+            path: The path of the file to open.
             mode: The mode to open the file.
 
         Returns:

zenml/artifacts/utils.py

@@ -575,7 +575,7 @@ def download_artifact_files_from_response(
                     )
                     file_path = str(Path(artifact.uri) / file_str)
                     with artifact_store.open(
-                        name=file_path, mode="rb"
+                        file_path, mode="rb"
                     ) as store_file:
                         # Use a loop to read and write chunks of the file
                         # instead of reading the entire file into memory

zenml/cli/__init__.py

@@ -2242,6 +2242,9 @@ export ZENML_STORE_URL=https://...
 export ZENML_STORE_API_KEY=<API_KEY>
 ```
 
+You don't need to run `zenml login` after setting these two environment
+variables and can start interacting with your server right away.
+
 To see all the service accounts you've created and their API keys, use the
 following commands:
 

zenml/cli/login.py

@@ -14,6 +14,7 @@
 """CLI for managing ZenML server deployments."""
 
 import ipaddress
+import os
 import re
 import sys
 import time
@@ -413,6 +414,27 @@ def connect_to_pro_server(
     cli_utils.declare(f"Connected to ZenML Pro server: {server.name}.")
 
 
+def _fail_if_authentication_environment_variables_set() -> None:
+    """Fail if any of the authentication environment variables are set."""
+    environment_variables = [
+        "ZENML_STORE_URL",
+        "ZENML_STORE_API_KEY",
+        "ZENML_STORE_USERNAME",
+        "ZENML_STORE_PASSWORD",
+    ]
+
+    if any(env_var in os.environ for env_var in environment_variables):
+        cli_utils.error(
+            "You're running to login/logout while having one of the "
+            f"{environment_variables} environment variables set. "
+            "If you want to use those environment variables to authenticate "
+            "to your ZenML server, there is no need to login/logout, you can "
+            "start interacting with your server right away. If you want to use "
+            "the `zenml login` command for authentication, please unset these "
+            "environment variables first."
+        )
+
+
 @cli.command(
     "login",
     help=(
@@ -670,6 +692,8 @@ def login(
             dashboard on a public domain. Primarily used for accessing the
             dashboard in Colab.
     """
+    _fail_if_authentication_environment_variables_set()
+
     if local:
         if api_key:
             cli_utils.error(
@@ -849,6 +873,8 @@ def logout(
     """
     from zenml.login.credentials_store import get_credentials_store
 
+    _fail_if_authentication_environment_variables_set()
+
     credentials_store = get_credentials_store()
     gc = GlobalConfiguration()
     store_cfg = gc.store_configuration

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/kubernetes/orchestrators/kube_utils.py

@@ -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:

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -395,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
 

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

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),
         }
     )
 

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",
+]

zenml/integrations/modal/flavors/modal_step_operator_flavor.py

@@ -0,0 +1,125 @@
+#  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 step operator flavor."""
+
+from typing import TYPE_CHECKING, Optional, Type
+
+from zenml.config.base_settings import BaseSettings
+from zenml.integrations.modal import MODAL_STEP_OPERATOR_FLAVOR
+from zenml.step_operators import BaseStepOperatorConfig, BaseStepOperatorFlavor
+
+if TYPE_CHECKING:
+    from zenml.integrations.modal.step_operators import ModalStepOperator
+
+
+class ModalStepOperatorSettings(BaseSettings):
+    """Settings for the Modal step operator.
+
+    Specifying the region and cloud provider is only available for Enterprise
+    and Team plan customers.
+
+    Certain combinations of settings are not available. It is suggested to err
+    on the side of looser settings rather than more restrictive ones to avoid
+    pipeline execution failures. In the case of failures, however, Modal
+    provides detailed error messages that can help identify what is
+    incompatible. See more in the Modal docs at https://modal.com/docs/guide/region-selection.
+
+    Attributes:
+        gpu: The type of GPU to use for the step execution.
+        region: The region to use for the step execution.
+        cloud: The cloud provider to use for the step execution.
+    """
+
+    gpu: Optional[str] = None
+    region: Optional[str] = None
+    cloud: Optional[str] = None
+
+
+class ModalStepOperatorConfig(
+    BaseStepOperatorConfig, ModalStepOperatorSettings
+):
+    """Configuration for the Modal step operator."""
+
+    @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
+
+
+class ModalStepOperatorFlavor(BaseStepOperatorFlavor):
+    """Modal step operator flavor."""
+
+    @property
+    def name(self) -> str:
+        """Name of the flavor.
+
+        Returns:
+            The name of the flavor.
+        """
+        return MODAL_STEP_OPERATOR_FLAVOR
+
+    @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/modal.png"
+
+    @property
+    def config_class(self) -> Type[ModalStepOperatorConfig]:
+        """Returns `ModalStepOperatorConfig` config class.
+
+        Returns:
+            The config class.
+        """
+        return ModalStepOperatorConfig
+
+    @property
+    def implementation_class(self) -> Type["ModalStepOperator"]:
+        """Implementation class for this flavor.
+
+        Returns:
+            The implementation class.
+        """
+        from zenml.integrations.modal.step_operators import ModalStepOperator
+
+        return ModalStepOperator

zenml/integrations/modal/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.
+"""Modal step operator."""
+
+from zenml.integrations.modal.step_operators.modal_step_operator import (
+    ModalStepOperator,
+)
+
+__all__ = [
+    "ModalStepOperator",
+]