zenml-nightly 0.70.0.dev20241127__py3-none-any.whl → 0.70.0.dev20241129__py3-none-any.whl

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

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/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: 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."

zenml/model/model.py

@@ -57,7 +57,9 @@ class Model(BaseModel):
     ethics: The ethical implications of the model.
     tags: Tags associated with the model.
     version: The version name, version number or stage is optional and points model context
-        to a specific version/stage. If skipped new version will be created.
+        to a specific version/stage. If skipped new version will be created. `version`
+        also supports placeholders: standard `{date}` and `{time}` and any custom placeholders
+        that are passed as substitutions in the pipeline or step decorators.
     save_models_to_registry: Whether to save all ModelArtifacts to Model Registry,
         if available in active stack.
     """
@@ -534,6 +536,8 @@ class Model(BaseModel):
         from zenml.models import ModelRequest
 
         zenml_client = Client()
+        # backup logic, if the Model class is used directly from the code
+        self.name = format_name_template(self.name, substitutions={})
         if self.model_version_id:
             mv = zenml_client.get_model_version(
                 model_version_name_or_number_or_id=self.model_version_id,
@@ -663,7 +667,7 @@ class Model(BaseModel):
 
         # backup logic, if the Model class is used directly from the code
         if isinstance(self.version, str):
-            self.version = format_name_template(self.version)
+            self.version = format_name_template(self.version, substitutions={})
 
         try:
             if self.version or self.model_version_id:

zenml/models/v2/core/pipeline_run.py

@@ -237,6 +237,10 @@ class PipelineRunResponseMetadata(WorkspaceScopedResponseMetadata):
         default=False,
         description="Whether a template can be created from this run.",
     )
+    steps_substitutions: Dict[str, Dict[str, str]] = Field(
+        title="Substitutions used in the step runs of this pipeline run.",
+        default_factory=dict,
+    )
 
 
 class PipelineRunResponseResources(WorkspaceScopedResponseResources):

zenml/models/v2/core/step_run.py

@@ -142,7 +142,7 @@ class StepRunRequest(WorkspaceScopedRequest):
 class StepRunUpdate(BaseModel):
     """Update model for step runs."""
 
-    outputs: Dict[str, UUID] = Field(
+    outputs: Dict[str, List[UUID]] = Field(
         title="The IDs of the output artifact versions of the step run.",
         default={},
     )

zenml/orchestrators/publish_utils.py

@@ -32,7 +32,7 @@ if TYPE_CHECKING:
 
 
 def publish_successful_step_run(
-    step_run_id: "UUID", output_artifact_ids: Dict[str, "UUID"]
+    step_run_id: "UUID", output_artifact_ids: Dict[str, List["UUID"]]
 ) -> "StepRunResponse":
     """Publishes a successful step run.
 

zenml/orchestrators/step_launcher.py

@@ -309,8 +309,12 @@ class StepLauncher:
             The created or existing pipeline run,
             and a boolean indicating whether the run was created or reused.
         """
+        start_time = datetime.utcnow()
         run_name = orchestrator_utils.get_run_name(
-            run_name_template=self._deployment.run_name_template
+            run_name_template=self._deployment.run_name_template,
+            substitutions=self._deployment.pipeline_configuration._get_full_substitutions(
+                start_time
+            ),
         )
 
         logger.debug("Creating pipeline run %s", run_name)
@@ -329,7 +333,7 @@ class StepLauncher:
             ),
             status=ExecutionStatus.RUNNING,
             orchestrator_environment=get_run_environment_dict(),
-            start_time=datetime.utcnow(),
+            start_time=start_time,
             tags=self._deployment.pipeline_configuration.tags,
         )
         return client.zen_store.get_or_create_run(pipeline_run)

zenml/orchestrators/step_run_utils.py

@@ -354,13 +354,16 @@ def create_cached_step_runs(
 
 
 def get_or_create_model_version_for_pipeline_run(
-    model: "Model", pipeline_run: PipelineRunResponse
+    model: "Model",
+    pipeline_run: PipelineRunResponse,
+    substitutions: Dict[str, str],
 ) -> Tuple[ModelVersionResponse, bool]:
     """Get or create a model version as part of a pipeline run.
 
     Args:
         model: The model to get or create.
         pipeline_run: The pipeline run for which the model should be created.
+        substitutions: Substitutions to apply to the model version name.
 
     Returns:
         The model version and a boolean indicating whether it was newly created
@@ -374,12 +377,14 @@ def get_or_create_model_version_for_pipeline_run(
         return model._get_model_version(), False
     elif model.version:
         if isinstance(model.version, str):
-            start_time = pipeline_run.start_time or datetime.utcnow()
             model.version = string_utils.format_name_template(
                 model.version,
-                date=start_time.strftime("%Y_%m_%d"),
-                time=start_time.strftime("%H_%M_%S_%f"),
+                substitutions=substitutions,
             )
+        model.name = string_utils.format_name_template(
+            model.name,
+            substitutions=substitutions,
+        )
 
         return (
             model._get_or_create_model_version(),
@@ -460,7 +465,9 @@ def prepare_pipeline_run_model_version(
         model_version = pipeline_run.model_version
     elif config_model := pipeline_run.config.model:
         model_version, _ = get_or_create_model_version_for_pipeline_run(
-            model=config_model, pipeline_run=pipeline_run
+            model=config_model,
+            pipeline_run=pipeline_run,
+            substitutions=pipeline_run.config.substitutions,
         )
         pipeline_run = Client().zen_store.update_run(
             run_id=pipeline_run.id,
@@ -492,7 +499,9 @@ def prepare_step_run_model_version(
         model_version = step_run.model_version
     elif config_model := step_run.config.model:
         model_version, created = get_or_create_model_version_for_pipeline_run(
-            model=config_model, pipeline_run=pipeline_run
+            model=config_model,
+            pipeline_run=pipeline_run,
+            substitutions=step_run.config.substitutions,
         )
         step_run = Client().zen_store.update_run_step(
             step_run_id=step_run.id,

zenml/orchestrators/step_runner.py

@@ -152,6 +152,15 @@ class StepRunner:
                 func=step_instance.entrypoint
             )
 
+            self._evaluate_artifact_names_in_collections(
+                step_run,
+                output_annotations,
+                [
+                    output_artifact_uris,
+                    output_materializers,
+                ],
+            )
+
             self._stack.prepare_step_run(info=step_run_info)
 
             # Initialize the step context singleton
@@ -257,7 +266,9 @@ class StepRunner:
 
             # Update the status and output artifacts of the step run.
             output_artifact_ids = {
-                output_name: artifact.id
+                output_name: [
+                    artifact.id,
+                ]
                 for output_name, artifact in output_artifacts.items()
             }
             publish_successful_step_run(
@@ -265,6 +276,33 @@ class StepRunner:
                 output_artifact_ids=output_artifact_ids,
             )
 
+    def _evaluate_artifact_names_in_collections(
+        self,
+        step_run: "StepRunResponse",
+        output_annotations: Dict[str, OutputSignature],
+        collections: List[Dict[str, Any]],
+    ) -> None:
+        """Evaluates the artifact names in the collections.
+
+        Args:
+            step_run: The step run.
+            output_annotations: The output annotations of the step function
+                (also evaluated).
+            collections: The collections to evaluate.
+        """
+        collections.append(output_annotations)
+        for k, v in list(output_annotations.items()):
+            _evaluated_name = None
+            if v.artifact_config:
+                _evaluated_name = v.artifact_config._evaluated_name(
+                    step_run.config.substitutions
+                )
+            if _evaluated_name is None:
+                _evaluated_name = k
+
+            for d in collections:
+                d[_evaluated_name] = d.pop(k)
+
     def _load_step(self) -> "BaseStep":
         """Load the step instance.
 
@@ -400,7 +438,7 @@ class StepRunner:
                 **artifact.get_hydrated_version().model_dump()
             )
 
-        if data_type is Any or is_union(get_origin(data_type)):
+        if data_type in (None, Any) or is_union(get_origin(data_type)):
             # Entrypoint function does not define a specific type for the input,
             # we use the datatype of the stored artifact
             data_type = source_utils.load(artifact.data_type)