zenml-nightly 0.70.0.dev20241125__py3-none-any.whl → 0.71.0.dev20241220__py3-none-any.whl

zenml/integrations/modal/step_operators/modal_step_operator.py

@@ -0,0 +1,242 @@
+#  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 implementation."""
+
+import asyncio
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type, cast
+
+import modal
+from modal_proto import api_pb2
+
+from zenml.client import Client
+from zenml.config.build_configuration import BuildConfiguration
+from zenml.config.resource_settings import ByteUnit, ResourceSettings
+from zenml.enums import StackComponentType
+from zenml.integrations.modal.flavors import (
+    ModalStepOperatorConfig,
+    ModalStepOperatorSettings,
+)
+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__)
+
+MODAL_STEP_OPERATOR_DOCKER_IMAGE_KEY = "modal_step_operator"
+
+
+def get_gpu_values(
+    settings: ModalStepOperatorSettings, resource_settings: ResourceSettings
+) -> Optional[str]:
+    """Get the GPU values for the Modal step operator.
+
+    Args:
+        settings: The Modal step operator settings.
+        resource_settings: The resource settings.
+
+    Returns:
+        The GPU string if a count is specified, otherwise the GPU type.
+    """
+    if not settings.gpu:
+        return None
+    gpu_count = resource_settings.gpu_count
+    return f"{settings.gpu}:{gpu_count}" if gpu_count else settings.gpu
+
+
+class ModalStepOperator(BaseStepOperator):
+    """Step operator to run a step on Modal.
+
+    This class defines code that can set up a Modal environment and run
+    functions in it.
+    """
+
+    @property
+    def config(self) -> ModalStepOperatorConfig:
+        """Get the Modal step operator configuration.
+
+        Returns:
+            The Modal step operator configuration.
+        """
+        return cast(ModalStepOperatorConfig, self._config)
+
+    @property
+    def settings_class(self) -> Optional[Type["BaseSettings"]]:
+        """Get the settings class for the Modal step operator.
+
+        Returns:
+            The Modal step operator settings class.
+        """
+        return ModalStepOperatorSettings
+
+    @property
+    def validator(self) -> Optional[StackValidator]:
+        """Get the stack validator for the Modal step operator.
+
+        Returns:
+            The stack validator.
+        """
+
+        def _validate_remote_components(stack: "Stack") -> Tuple[bool, str]:
+            if stack.artifact_store.config.is_local:
+                return False, (
+                    "The Modal 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 Modal "
+                    "step operator."
+                )
+
+            container_registry = stack.container_registry
+            assert container_registry is not None
+
+            if container_registry.config.is_local:
+                return False, (
+                    "The Modal 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 "
+                    "Modal 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"]:
+        """Get the Docker build configurations for the Modal step operator.
+
+        Args:
+            deployment: The pipeline deployment.
+
+        Returns:
+            A list of Docker build configurations.
+        """
+        builds = []
+        for step_name, step in deployment.step_configurations.items():
+            if step.config.step_operator == self.name:
+                build = BuildConfiguration(
+                    key=MODAL_STEP_OPERATOR_DOCKER_IMAGE_KEY,
+                    settings=step.config.docker_settings,
+                    step_name=step_name,
+                )
+                builds.append(build)
+
+        return builds
+
+    def launch(
+        self,
+        info: "StepRunInfo",
+        entrypoint_command: List[str],
+        environment: Dict[str, str],
+    ) -> None:
+        """Launch a step run on Modal.
+
+        Args:
+            info: The step run information.
+            entrypoint_command: The entrypoint command for the step.
+            environment: The environment variables for the step.
+
+        Raises:
+            RuntimeError: If no Docker credentials are found for the container registry.
+            ValueError: If no container registry is found in the stack.
+        """
+        settings = cast(ModalStepOperatorSettings, self.get_settings(info))
+        image_name = info.get_image(key=MODAL_STEP_OPERATOR_DOCKER_IMAGE_KEY)
+        zc = Client()
+        stack = zc.active_stack
+
+        if not stack.container_registry:
+            raise ValueError(
+                "No Container registry found in the stack. "
+                "Please add a container registry and ensure "
+                "it is correctly configured."
+            )
+
+        if docker_creds := stack.container_registry.credentials:
+            docker_username, docker_password = docker_creds
+        else:
+            raise RuntimeError(
+                "No Docker credentials found for the container registry."
+            )
+
+        my_secret = modal.secret._Secret.from_dict(
+            {
+                "REGISTRY_USERNAME": docker_username,
+                "REGISTRY_PASSWORD": docker_password,
+            }
+        )
+
+        spec = modal.image.DockerfileSpec(
+            commands=[f"FROM {image_name}"], context_files={}
+        )
+
+        zenml_image = modal.Image._from_args(
+            dockerfile_function=lambda *_, **__: spec,
+            force_build=False,
+            image_registry_config=modal.image._ImageRegistryConfig(
+                api_pb2.REGISTRY_AUTH_TYPE_STATIC_CREDS, my_secret
+            ),
+        ).env(environment)
+
+        resource_settings = info.config.resource_settings
+        gpu_values = get_gpu_values(settings, resource_settings)
+
+        app = modal.App(
+            f"zenml-{info.run_name}-{info.step_run_id}-{info.pipeline_step_name}"
+        )
+
+        async def run_sandbox() -> asyncio.Future[None]:
+            loop = asyncio.get_event_loop()
+            future = loop.create_future()
+            with modal.enable_output():
+                async with app.run():
+                    memory_mb = resource_settings.get_memory(ByteUnit.MB)
+                    memory_int = (
+                        int(memory_mb) if memory_mb is not None else None
+                    )
+                    sb = await modal.Sandbox.create.aio(
+                        "bash",
+                        "-c",
+                        " ".join(entrypoint_command),
+                        image=zenml_image,
+                        gpu=gpu_values,
+                        cpu=resource_settings.cpu_count,
+                        memory=memory_int,
+                        cloud=settings.cloud,
+                        region=settings.region,
+                        app=app,
+                        timeout=86400,  # 24h, the max Modal allows
+                    )
+
+                    await sb.wait.aio()
+
+            future.set_result(None)
+            return future
+
+        asyncio.run(run_sandbox())

zenml/integrations/neptune/experiment_trackers/neptune_experiment_tracker.py

@@ -77,10 +77,12 @@ class NeptuneExperimentTracker(BaseExperimentTracker):
             NeptuneExperimentTrackerSettings, self.get_settings(info)
         )
 
-        self.run_state.token = self.config.api_token
-        self.run_state.project = self.config.project
-        self.run_state.run_name = info.run_name
-        self.run_state.tags = list(settings.tags)
+        self.run_state.initialize(
+            project=self.config.project,
+            token=self.config.api_token,
+            run_name=info.run_name,
+            tags=list(settings.tags),
+        )
 
     def get_step_run_metadata(
         self, info: "StepRunInfo"
@@ -107,4 +109,4 @@ class NeptuneExperimentTracker(BaseExperimentTracker):
         """
         self.run_state.active_run.sync()
         self.run_state.active_run.stop()
-        self.run_state.reset_active_run()
+        self.run_state.reset()

zenml/integrations/neptune/experiment_trackers/run_state.py

@@ -20,7 +20,6 @@ import neptune
 
 import zenml
 from zenml.client import Client
-from zenml.integrations.constants import NEPTUNE
 from zenml.utils.singleton import SingletonMetaClass
 
 if TYPE_CHECKING:
@@ -29,20 +28,38 @@ if TYPE_CHECKING:
 _INTEGRATION_VERSION_KEY = "source_code/integrations/zenml"
 
 
-class InvalidExperimentTrackerSelected(Exception):
-    """Raised if a Neptune run is fetched while using a different experiment tracker."""
-
-
 class RunProvider(metaclass=SingletonMetaClass):
     """Singleton object used to store and persist a Neptune run state across the pipeline."""
 
     def __init__(self) -> None:
         """Initialize RunProvider. Called with no arguments."""
         self._active_run: Optional["Run"] = None
-        self._project: Optional[str]
-        self._run_name: Optional[str]
-        self._token: Optional[str]
-        self._tags: Optional[List[str]]
+        self._project: Optional[str] = None
+        self._run_name: Optional[str] = None
+        self._token: Optional[str] = None
+        self._tags: Optional[List[str]] = None
+        self._initialized = False
+
+    def initialize(
+        self,
+        project: Optional[str] = None,
+        token: Optional[str] = None,
+        run_name: Optional[str] = None,
+        tags: Optional[List[str]] = None,
+    ) -> None:
+        """Initialize the run state.
+
+        Args:
+            project: The neptune project.
+            token: The neptune token.
+            run_name: The neptune run name.
+            tags: Tags for the neptune run.
+        """
+        self._project = project
+        self._token = token
+        self._run_name = run_name
+        self._tags = tags
+        self._initialized = True
 
     @property
     def project(self) -> Optional[Any]:
@@ -53,15 +70,6 @@ class RunProvider(metaclass=SingletonMetaClass):
         """
         return self._project
 
-    @project.setter
-    def project(self, project: str) -> None:
-        """Setter for project name.
-
-        Args:
-            project: Neptune project name
-        """
-        self._project = project
-
     @property
     def token(self) -> Optional[Any]:
         """Getter for API token.
@@ -71,15 +79,6 @@ class RunProvider(metaclass=SingletonMetaClass):
         """
         return self._token
 
-    @token.setter
-    def token(self, token: str) -> None:
-        """Setter for API token.
-
-        Args:
-            token: Neptune API token
-        """
-        self._token = token
-
     @property
     def run_name(self) -> Optional[Any]:
         """Getter for run name.
@@ -89,15 +88,6 @@ class RunProvider(metaclass=SingletonMetaClass):
         """
         return self._run_name
 
-    @run_name.setter
-    def run_name(self, run_name: str) -> None:
-        """Setter for run name.
-
-        Args:
-            run_name: name of the pipeline run
-        """
-        self._run_name = run_name
-
     @property
     def tags(self) -> Optional[Any]:
         """Getter for run tags.
@@ -107,14 +97,14 @@ class RunProvider(metaclass=SingletonMetaClass):
         """
         return self._tags
 
-    @tags.setter
-    def tags(self, tags: List[str]) -> None:
-        """Setter for run tags.
+    @property
+    def initialized(self) -> bool:
+        """If the run state is initialized.
 
-        Args:
-            tags: list of tags associated with a Neptune run
+        Returns:
+            If the run state is initialized.
         """
-        self._tags = tags
+        return self._initialized
 
     @property
     def active_run(self) -> "Run":
@@ -137,9 +127,14 @@ class RunProvider(metaclass=SingletonMetaClass):
             self._active_run = run
         return self._active_run
 
-    def reset_active_run(self) -> None:
-        """Resets the active run state to None."""
+    def reset(self) -> None:
+        """Reset the run state."""
         self._active_run = None
+        self._project = None
+        self._run_name = None
+        self._token = None
+        self._tags = None
+        self._initialized = False
 
 
 def get_neptune_run() -> "Run":
@@ -149,14 +144,35 @@ def get_neptune_run() -> "Run":
         Neptune run object
 
     Raises:
-        InvalidExperimentTrackerSelected: when called while using an experiment tracker other than Neptune
+        RuntimeError: When unable to fetch the active neptune run.
     """
-    client = Client()
-    experiment_tracker = client.active_stack.experiment_tracker
-    if experiment_tracker.flavor == NEPTUNE:  # type: ignore
-        return experiment_tracker.run_state.active_run  # type: ignore
-    raise InvalidExperimentTrackerSelected(
-        "Fetching a Neptune run works only with the 'neptune' flavor of "
-        "the experiment tracker. The flavor currently selected is %s"
-        % experiment_tracker.flavor  # type: ignore
+    from zenml.integrations.neptune.experiment_trackers import (
+        NeptuneExperimentTracker,
     )
+
+    experiment_tracker = Client().active_stack.experiment_tracker
+
+    if not experiment_tracker:
+        raise RuntimeError(
+            "Unable to get neptune run: Missing experiment tracker in the "
+            "active stack."
+        )
+
+    if not isinstance(experiment_tracker, NeptuneExperimentTracker):
+        raise RuntimeError(
+            "Unable to get neptune run: Experiment tracker in the active "
+            f"stack ({experiment_tracker.flavor}) is not a neptune experiment "
+            "tracker."
+        )
+
+    run_state = experiment_tracker.run_state
+    if not run_state.initialized:
+        raise RuntimeError(
+            "Unable to get neptune run: The experiment tracker has not been "
+            "initialized. To solve this, make sure you use the experiment "
+            "tracker in your step. See "
+            "https://docs.zenml.io/stack-components/experiment-trackers/neptune#how-do-you-use-it "
+            "for more information."
+        )
+
+    return experiment_tracker.run_state.active_run

zenml/integrations/registry.py

@@ -111,7 +111,7 @@ class IntegrationRegistry(object):
                 )
             else:
                 raise KeyError(
-                    f"Version {integration_name} does not exist. "
+                    f"Integration {integration_name} does not exist. "
                     f"Currently the following integrations are implemented. "
                     f"{self.list_integration_names}"
                 )
@@ -148,7 +148,7 @@ class IntegrationRegistry(object):
                 ].get_uninstall_requirements(target_os=target_os)
             else:
                 raise KeyError(
-                    f"Version {integration_name} does not exist. "
+                    f"Integration {integration_name} does not exist. "
                     f"Currently the following integrations are implemented. "
                     f"{self.list_integration_names}"
                 )

zenml/integrations/skypilot/flavors/skypilot_orchestrator_base_vm_config.py

@@ -144,3 +144,15 @@ class SkypilotBaseOrchestratorConfig(
             True if this config is for a local component, False otherwise.
         """
         return False
+
+    @property
+    def supports_client_side_caching(self) -> bool:
+        """Whether the orchestrator supports client side caching.
+
+        Returns:
+            Whether the orchestrator supports client side caching.
+        """
+        # The Skypilot orchestrator runs the entire pipeline in a single VM, or
+        # starts additional VMs from the root VM. Both of those cases are
+        # currently not supported when using client-side caching.
+        return False

zenml/integrations/wandb/flavors/wandb_experiment_tracker_flavor.py

@@ -23,7 +23,7 @@ from typing import (
     cast,
 )
 
-from pydantic import field_validator
+from pydantic import field_validator, BaseModel
 
 from zenml.config.base_settings import BaseSettings
 from zenml.experiment_trackers.base_experiment_tracker import (
@@ -60,18 +60,26 @@ class WandbExperimentTrackerSettings(BaseSettings):
         Args:
             value: The settings.
 
+        Raises:
+            ValueError: If converting the settings failed.
+
         Returns:
             Dict representation of the settings.
         """
         import wandb
 
         if isinstance(value, wandb.Settings):
-            # Depending on the wandb version, either `make_static` or `to_dict`
-            # is available to convert the settings to a dictionary
-            if hasattr(value, "make_static"):
+            # Depending on the wandb version, either `model_dump`, 
+            # `make_static` or `to_dict` is available to convert the settings 
+            # to a dictionary
+            if isinstance(value, BaseModel):
+                return value.model_dump()
+            elif hasattr(value, "make_static"):
                 return cast(Dict[str, Any], value.make_static())
-            else:
+            elif hasattr(value, "to_dict"):
                 return value.to_dict()
+            else:
+                raise ValueError("Unable to convert wandb settings to dict.")
         else:
             return value
 

zenml/io/filesystem.py

@@ -54,11 +54,11 @@ class BaseFilesystem(ABC):
 
     @staticmethod
     @abstractmethod
-    def open(name: PathType, mode: str = "r") -> Any:
+    def open(path: PathType, mode: str = "r") -> Any:
         """Opens a file.
 
         Args:
-            name: The path to the file.
+            path: The path to the file.
             mode: The mode to open the file in.
 
         Returns:

zenml/io/local_filesystem.py

@@ -55,18 +55,18 @@ class LocalFilesystem(BaseFilesystem):
     SUPPORTED_SCHEMES: ClassVar[Set[str]] = {""}
 
     @staticmethod
-    def open(name: PathType, mode: str = "r") -> Any:
+    def open(path: PathType, mode: str = "r") -> Any:
         """Open a file at the given path.
 
         Args:
-            name: The path to the file.
+            path: The path to the file.
             mode: The mode to open the file.
 
         Returns:
             Any: The file object.
         """
         encoding = "utf-8" if "b" not in mode else None
-        return open(name, mode=mode, encoding=encoding)
+        return open(path, mode=mode, encoding=encoding)
 
     @staticmethod
     def copyfile(

zenml/materializers/built_in_materializer.py

@@ -28,7 +28,7 @@ from typing import (
 )
 
 from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
-from zenml.enums import ArtifactType
+from zenml.enums import ArtifactType, VisualizationType
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.materializers.materializer_registry import materializer_registry
@@ -415,6 +415,23 @@ class BuiltInContainerMaterializer(BaseMaterializer):
                 self.artifact_store.rmtree(entry["path"])
             raise e
 
+    # save dict type objects to JSON file with JSON visualization type
+    def save_visualizations(self, data: Any) -> Dict[str, "VisualizationType"]:
+        """Save visualizations for the given data.
+
+        Args:
+            data: The data to save visualizations for.
+
+        Returns:
+            A dictionary of visualization URIs and their types.
+        """
+        # dict/list type objects are always saved as JSON files
+        # doesn't work for non-serializable types as they
+        # are saved as list of lists in different files
+        if _is_serializable(data):
+            return {self.data_path.replace("\\", "/"): VisualizationType.JSON}
+        return {}
+
     def extract_metadata(self, data: Any) -> Dict[str, "MetadataType"]:
         """Extract metadata from the given built-in container object.
 

zenml/materializers/structured_string_materializer.py

@@ -19,22 +19,23 @@ from typing import Dict, Type, Union
 from zenml.enums import ArtifactType, VisualizationType
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
-from zenml.types import CSVString, HTMLString, MarkdownString
+from zenml.types import CSVString, HTMLString, JSONString, MarkdownString
 
 logger = get_logger(__name__)
 
 
-STRUCTURED_STRINGS = Union[CSVString, HTMLString, MarkdownString]
+STRUCTURED_STRINGS = Union[CSVString, HTMLString, MarkdownString, JSONString]
 
 HTML_FILENAME = "output.html"
 MARKDOWN_FILENAME = "output.md"
 CSV_FILENAME = "output.csv"
+JSON_FILENAME = "output.json"
 
 
 class StructuredStringMaterializer(BaseMaterializer):
     """Materializer for HTML or Markdown strings."""
 
-    ASSOCIATED_TYPES = (CSVString, HTMLString, MarkdownString)
+    ASSOCIATED_TYPES = (CSVString, HTMLString, MarkdownString, JSONString)
     ASSOCIATED_ARTIFACT_TYPE = ArtifactType.DATA_ANALYSIS
 
     def load(self, data_type: Type[STRUCTURED_STRINGS]) -> STRUCTURED_STRINGS:
@@ -94,6 +95,8 @@ class StructuredStringMaterializer(BaseMaterializer):
             filename = HTML_FILENAME
         elif issubclass(data_type, MarkdownString):
             filename = MARKDOWN_FILENAME
+        elif issubclass(data_type, JSONString):
+            filename = JSON_FILENAME
         else:
             raise ValueError(
                 f"Data type {data_type} is not supported by this materializer."
@@ -120,6 +123,8 @@ class StructuredStringMaterializer(BaseMaterializer):
             return VisualizationType.HTML
         elif issubclass(data_type, MarkdownString):
             return VisualizationType.MARKDOWN
+        elif issubclass(data_type, JSONString):
+            return VisualizationType.JSON
         else:
             raise ValueError(
                 f"Data type {data_type} is not supported by this materializer."