zenml-nightly 0.70.0.dev20241126__py3-none-any.whl → 0.70.0.dev20241128__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/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/model/model.py

@@ -13,14 +13,12 @@
 #  permissions and limitations under the License.
 """Model user facing interface to pass into pipeline or step."""
 
-import datetime
 from typing import (
     TYPE_CHECKING,
     Any,
     Dict,
     List,
     Optional,
-    Tuple,
     Union,
 )
 from uuid import UUID
@@ -41,7 +39,6 @@ if TYPE_CHECKING:
         ModelResponse,
         ModelVersionResponse,
         PipelineRunResponse,
-        StepRunResponse,
     )
 
 logger = get_logger(__name__)
@@ -743,85 +740,6 @@ class Model(BaseModel):
             )
         )
 
-    def _prepare_model_version_before_step_launch(
-        self,
-        pipeline_run: "PipelineRunResponse",
-        step_run: Optional["StepRunResponse"],
-        return_logs: bool,
-    ) -> Tuple[str, "PipelineRunResponse", Optional["StepRunResponse"]]:
-        """Prepares model version inside pipeline run.
-
-        Args:
-            pipeline_run: pipeline run
-            step_run: step run (passed only if model version is defined in a step explicitly)
-            return_logs: whether to return logs or not
-
-        Returns:
-            Logs related to the Dashboard URL to show later.
-        """
-        from zenml.client import Client
-        from zenml.models import PipelineRunUpdate, StepRunUpdate
-
-        logs = ""
-
-        # copy Model instance to prevent corrupting configs of the
-        # subsequent runs, if they share the same config object
-        self_copy = self.model_copy()
-
-        # in case request is within the step and no self-configuration is provided
-        # try reuse what's in the pipeline run first
-        if step_run is None and pipeline_run.model_version is not None:
-            self_copy.version = pipeline_run.model_version.name
-            self_copy.model_version_id = pipeline_run.model_version.id
-        # otherwise try to fill the templated name, if needed
-        elif isinstance(self_copy.version, str):
-            if pipeline_run.start_time:
-                start_time = pipeline_run.start_time
-            else:
-                start_time = datetime.datetime.now(datetime.timezone.utc)
-            self_copy.version = format_name_template(
-                self_copy.version,
-                date=start_time.strftime("%Y_%m_%d"),
-                time=start_time.strftime("%H_%M_%S_%f"),
-            )
-
-        # if exact model not yet defined - try to get/create and update it
-        # back to the run accordingly
-        if self_copy.model_version_id is None:
-            model_version_response = self_copy._get_or_create_model_version()
-
-            client = Client()
-            # update the configured model version id in runs accordingly
-            if step_run:
-                step_run = client.zen_store.update_run_step(
-                    step_run_id=step_run.id,
-                    step_run_update=StepRunUpdate(
-                        model_version_id=model_version_response.id
-                    ),
-                )
-            else:
-                pipeline_run = client.zen_store.update_run(
-                    run_id=pipeline_run.id,
-                    run_update=PipelineRunUpdate(
-                        model_version_id=model_version_response.id
-                    ),
-                )
-
-            if return_logs:
-                from zenml.utils.cloud_utils import try_get_model_version_url
-
-                if logs_to_show := try_get_model_version_url(
-                    model_version_response
-                ):
-                    logs = logs_to_show
-                else:
-                    logs = (
-                        "Models can be viewed in the dashboard using ZenML Pro. Sign up "
-                        "for a free trial at https://www.zenml.io/pro/"
-                    )
-        self.model_version_id = self_copy.model_version_id
-        return logs, pipeline_run, step_run
-
     @property
     def _lazy_version(self) -> Optional[str]:
         """Get version name for lazy loader.

zenml/orchestrators/step_run_utils.py

@@ -518,10 +518,15 @@ def log_model_version_dashboard_url(
     Args:
         model_version: The model version for which to log the dashboard URL.
     """
-    from zenml.utils.cloud_utils import try_get_model_version_url
+    from zenml.utils.dashboard_utils import get_model_version_url
 
-    if model_version_url_logs := try_get_model_version_url(model_version):
-        logger.info(model_version_url_logs)
+    if model_version_url := get_model_version_url(model_version.id):
+        logger.info(
+            "Dashboard URL for Model Version `%s (%s)`:\n%s",
+            model_version.model.name,
+            model_version.name,
+            model_version_url,
+        )
     else:
         logger.info(
             "Models can be viewed in the dashboard using ZenML Pro. Sign up "

zenml/orchestrators/step_runner.py

@@ -400,7 +400,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)

zenml/orchestrators/utils.py

@@ -40,7 +40,9 @@ if TYPE_CHECKING:
     from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
 
 
-def get_orchestrator_run_name(pipeline_name: str) -> str:
+def get_orchestrator_run_name(
+    pipeline_name: str, max_length: Optional[int] = None
+) -> str:
     """Gets an orchestrator run name.
 
     This run name is not the same as the ZenML run name but can instead be
@@ -48,11 +50,31 @@ def get_orchestrator_run_name(pipeline_name: str) -> str:
 
     Args:
         pipeline_name: Name of the pipeline that will run.
+        max_length: Maximum length of the generated name.
+
+    Raises:
+        ValueError: If the max length is below 8 characters.
 
     Returns:
         The orchestrator run name.
     """
-    return f"{pipeline_name}_{random.Random().getrandbits(128):032x}"
+    suffix_length = 32
+    pipeline_name = f"{pipeline_name}_"
+
+    if max_length:
+        if max_length < 8:
+            raise ValueError(
+                "Maximum length for orchestrator run name must be 8 or above."
+            )
+
+        # Make sure we always have a certain suffix to guarantee no overlap
+        # with other runs
+        suffix_length = min(32, max(8, max_length - len(pipeline_name)))
+        pipeline_name = pipeline_name[: (max_length - suffix_length)]
+
+    suffix = "".join(random.choices("0123456789abcdef", k=suffix_length))
+
+    return f"{pipeline_name}{suffix}"
 
 
 def is_setting_enabled(

zenml/steps/entrypoint_function_utils.py

@@ -195,7 +195,9 @@ class EntrypointFunctionDefinition(NamedTuple):
             parameter: The function parameter for which the value was provided.
             value: The input value.
         """
-        config_dict = ConfigDict(arbitrary_types_allowed=False)
+        # We allow passing None for optional annotations that would otherwise
+        # not be allowed as a parameter
+        config_dict = ConfigDict(arbitrary_types_allowed=value is None)
 
         # Create a pydantic model with just a single required field with the
         # type annotation of the parameter to verify the input type including

zenml/zen_server/cloud_utils.py

@@ -170,7 +170,9 @@ class ZenMLCloudConnection:
             token = self._fetch_auth_token()
             self._session.headers.update({"Authorization": "Bearer " + token})
 
-            retries = Retry(total=5, backoff_factor=0.1)
+            retries = Retry(
+                total=5, backoff_factor=0.1, status_forcelist=[502, 504]
+            )
             self._session.mount(
                 "https://",
                 HTTPAdapter(

zenml/zen_server/rbac/endpoint_utils.py

@@ -189,7 +189,7 @@ def verify_permissions_and_list_entities(
 def verify_permissions_and_update_entity(
     id: UUIDOrStr,
     update_model: AnyUpdate,
-    get_method: Callable[[UUIDOrStr], AnyResponse],
+    get_method: Callable[[UUIDOrStr, bool], AnyResponse],
     update_method: Callable[[UUIDOrStr, AnyUpdate], AnyResponse],
 ) -> AnyResponse:
     """Verify permissions and update an entity.
@@ -203,7 +203,8 @@ def verify_permissions_and_update_entity(
     Returns:
         A model of the updated entity.
     """
-    model = get_method(id)
+    # We don't need the hydrated version here
+    model = get_method(id, False)
     verify_permission_for_model(model, action=Action.UPDATE)
     updated_model = update_method(model.id, update_model)
     return dehydrate_response_model(updated_model)
@@ -211,7 +212,7 @@ def verify_permissions_and_update_entity(
 
 def verify_permissions_and_delete_entity(
     id: UUIDOrStr,
-    get_method: Callable[[UUIDOrStr], AnyResponse],
+    get_method: Callable[[UUIDOrStr, bool], AnyResponse],
     delete_method: Callable[[UUIDOrStr], None],
 ) -> AnyResponse:
     """Verify permissions and delete an entity.
@@ -224,7 +225,8 @@ def verify_permissions_and_delete_entity(
     Returns:
         The deleted entity.
     """
-    model = get_method(id)
+    # We don't need the hydrated version here
+    model = get_method(id, False)
     verify_permission_for_model(model, action=Action.DELETE)
     delete_method(model.id)
 

zenml/zen_server/rbac/models.py

@@ -59,7 +59,6 @@ class ResourceType(StrEnum):
     PIPELINE_DEPLOYMENT = "pipeline_deployment"
     PIPELINE_BUILD = "pipeline_build"
     RUN_TEMPLATE = "run_template"
-    USER = "user"
     SERVICE = "service"
     RUN_METADATA = "run_metadata"
     SECRET = "secret"
@@ -70,7 +69,9 @@ class ResourceType(StrEnum):
     TAG = "tag"
     TRIGGER = "trigger"
     TRIGGER_EXECUTION = "trigger_execution"
-    WORKSPACE = "workspace"
+    # Deactivated for now
+    # USER = "user"
+    # WORKSPACE = "workspace"
 
 
 class Resource(BaseModel):

zenml/zen_server/rbac/utils.py

@@ -413,8 +413,6 @@ def get_resource_type_for_model(
         TagResponse,
         TriggerExecutionResponse,
         TriggerResponse,
-        UserResponse,
-        WorkspaceResponse,
     )
 
     mapping: Dict[
@@ -434,8 +432,8 @@ def get_resource_type_for_model(
         ModelVersionResponse: ResourceType.MODEL_VERSION,
         ArtifactResponse: ResourceType.ARTIFACT,
         ArtifactVersionResponse: ResourceType.ARTIFACT_VERSION,
-        WorkspaceResponse: ResourceType.WORKSPACE,
-        UserResponse: ResourceType.USER,
+        # WorkspaceResponse: ResourceType.WORKSPACE,
+        # UserResponse: ResourceType.USER,
         PipelineDeploymentResponse: ResourceType.PIPELINE_DEPLOYMENT,
         PipelineBuildResponse: ResourceType.PIPELINE_BUILD,
         PipelineRunResponse: ResourceType.PIPELINE_RUN,
@@ -570,7 +568,6 @@ def get_schema_for_resource_type(
         TriggerExecutionSchema,
         TriggerSchema,
         UserSchema,
-        WorkspaceSchema,
     )
 
     mapping: Dict[ResourceType, Type["BaseSchema"]] = {
@@ -588,13 +585,13 @@ def get_schema_for_resource_type(
         ResourceType.SERVICE: ServiceSchema,
         ResourceType.TAG: TagSchema,
         ResourceType.SERVICE_ACCOUNT: UserSchema,
-        ResourceType.WORKSPACE: WorkspaceSchema,
+        # ResourceType.WORKSPACE: WorkspaceSchema,
         ResourceType.PIPELINE_RUN: PipelineRunSchema,
         ResourceType.PIPELINE_DEPLOYMENT: PipelineDeploymentSchema,
         ResourceType.PIPELINE_BUILD: PipelineBuildSchema,
         ResourceType.RUN_TEMPLATE: RunTemplateSchema,
         ResourceType.RUN_METADATA: RunMetadataSchema,
-        ResourceType.USER: UserSchema,
+        # ResourceType.USER: UserSchema,
         ResourceType.ACTION: ActionSchema,
         ResourceType.EVENT_SOURCE: EventSourceSchema,
         ResourceType.TRIGGER: TriggerSchema,