Follow up fixes for deployments (#4013)

Author: schustmi

src/zenml/artifact_stores/base_artifact_store.py

@@ -38,6 +38,7 @@
 from zenml.constants import (
     ENV_ZENML_SERVER,
     ENV_ZENML_SERVER_ALLOW_LOCAL_FILE_ACCESS,
+    IN_MEMORY_ARTIFACT_URI_PREFIX,
     handle_bool_env_var,
 )
 from zenml.enums import StackComponentType
@@ -106,8 +107,8 @@ def _validate_path(self, path: str) -> None:
             IllegalOperationError: If the path is a local file and the server
                 is not configured to allow local file access.
         """
-        # Skip validation for memory:// URIs used in serving mode
-        if path.startswith("memory://"):
+        if path.startswith(IN_MEMORY_ARTIFACT_URI_PREFIX):
+            # No need to validate in-memory URIs
             return
 
         if not self.allow_local_file_access and not io_utils.is_remote(path):
@@ -143,9 +144,7 @@ def _sanitize_potential_path(self, potential_path: Any) -> Any:
             # Neither string nor bytes, this is not a path
             return potential_path
 
-        # Preserve special in-memory scheme used by serving mode as-is
-        # to avoid treating it as a local filesystem path.
-        if isinstance(path, str) and path.startswith("memory://"):
+        if path.startswith(IN_MEMORY_ARTIFACT_URI_PREFIX):
             return path
 
         if io_utils.is_remote(path):

src/zenml/artifacts/utils.py

@@ -38,6 +38,7 @@
 from zenml.client import Client
 from zenml.constants import (
     ENV_ZENML_SERVER,
+    IN_MEMORY_ARTIFACT_URI_PREFIX,
     MODEL_METADATA_YAML_FILE_NAME,
 )
 from zenml.enums import (
@@ -54,6 +55,9 @@
 )
 from zenml.io import fileio
 from zenml.logger import get_logger
+from zenml.materializers.in_memory_materializer import (
+    InMemoryMaterializer,
+)
 from zenml.metadata.metadata_types import validate_metadata
 from zenml.models import (
     ArtifactVersionRequest,
@@ -153,20 +157,14 @@ def _store_artifact_data_and_prepare_request(
     """
     artifact_store = Client().active_stack.artifact_store
 
-    # Detect in-memory materializer to avoid touching the artifact store.
-    # Local import to minimize import-time dependencies.
-    from zenml.materializers.in_memory_materializer import (
-        InMemoryMaterializer,
+    uses_in_memory_materializer = issubclass(
+        materializer_class, InMemoryMaterializer
     )
 
-    is_in_memory = issubclass(materializer_class, InMemoryMaterializer)
-
-    if not is_in_memory:
+    if not uses_in_memory_materializer:
         artifact_store.makedirs(uri)
-    else:
-        # Ensure URI clearly indicates in-memory storage and not the artifact store
-        if not uri.startswith("memory://"):
-            uri = f"memory://custom_artifacts/{name}/{uuid4()}"
+    elif not uri.startswith(IN_MEMORY_ARTIFACT_URI_PREFIX):
+        uri = f"{IN_MEMORY_ARTIFACT_URI_PREFIX}{name}/{uuid4()}"
 
     materializer = materializer_class(uri=uri, artifact_store=artifact_store)
     materializer.uri = materializer.uri.replace("\\", "/")
@@ -204,7 +202,9 @@ def _store_artifact_data_and_prepare_request(
         data_type=source_utils.resolve(data_type),
         content_hash=content_hash,
         project=Client().active_project.id,
-        artifact_store_id=None if is_in_memory else artifact_store.id,
+        artifact_store_id=None
+        if uses_in_memory_materializer
+        else artifact_store.id,
         visualizations=visualizations,
         has_custom_name=has_custom_name,
         save_type=save_type,

src/zenml/cli/deployment.py

@@ -93,7 +93,7 @@ def list_deployments(**kwargs: Any) -> None:
         cli_utils.print_page_info(deployments)
 
 
-@deployment.command("describe")
+@deployment.command("describe", help="Describe a deployment.")
 @click.argument("deployment_name_or_id", type=str, required=True)
 @click.option(
     "--show-secret",
@@ -151,7 +151,7 @@ def describe_deployment(
         )
 
 
-@deployment.command("provision")
+@deployment.command("provision", help="Provision a deployment.")
 @click.argument("deployment_name_or_id", type=str, required=True)
 @click.option(
     "--snapshot",
@@ -253,7 +253,7 @@ def provision_deployment(
             cli_utils.pretty_print_deployment(deployment, show_secret=True)
 
 
-@deployment.command("deprovision")
+@deployment.command("deprovision", help="Deprovision a deployment.")
 @click.argument("deployment_name_or_id", type=str, required=False)
 @click.option(
     "--all",
@@ -408,7 +408,7 @@ def deprovision_deployment(
                     cli_utils.error(error_message)
 
 
-@deployment.command("delete")
+@deployment.command("delete", help="Delete a deployment.")
 @click.argument("deployment_name_or_id", type=str, required=False)
 @click.option(
     "--all",
@@ -564,7 +564,7 @@ def delete_deployment(
                     cli_utils.error(error_message)
 
 
-@deployment.command("refresh")
+@deployment.command("refresh", help="Refresh the status of a deployment.")
 @click.argument("deployment_name_or_id", type=str, required=True)
 def refresh_deployment(
     deployment_name_or_id: str,
@@ -586,7 +586,9 @@ def refresh_deployment(
 
 
 @deployment.command(
-    "invoke", context_settings={"ignore_unknown_options": True}
+    "invoke",
+    context_settings={"ignore_unknown_options": True},
+    help="Invoke a deployment with arguments.",
 )
 @click.argument("deployment_name_or_id", type=str, required=True)
 @click.option(
@@ -656,7 +658,7 @@ def invoke_deployment(
             )
 
 
-@deployment.command("logs")
+@deployment.command("logs", help="Get the logs of a deployment.")
 @click.argument("deployment_name_or_id", type=str, required=True)
 @click.option(
     "--follow",
@@ -694,8 +696,5 @@ def log_deployment(
     except KeyError as e:
         cli_utils.error(str(e))
     else:
-        with console.status(
-            f"Streaming logs for deployment '{deployment_name_or_id}'...\n"
-        ):
-            for log in logs:
-                print(log)
+        for log in logs:
+            print(log)

src/zenml/cli/utils.py

@@ -2383,7 +2383,7 @@ def get_deployment_status_emoji(
     if status == DeploymentStatus.ERROR:
         return ":x:"
     if status == DeploymentStatus.RUNNING:
-        return ":gear:"
+        return ":green_circle:"
     if status == DeploymentStatus.ABSENT:
         return ":stop_sign:"
 
@@ -2437,6 +2437,7 @@ def print_deployment_table(
         status = deployment.status or DeploymentStatus.UNKNOWN.value
         status_emoji = get_deployment_status_emoji(status)
         run_dict = {
+            "ID": deployment.id,
             "NAME": deployment.name,
             "PIPELINE": pipeline_name,
             "SNAPSHOT": deployment.snapshot.name or ""

src/zenml/constants.py

@@ -461,6 +461,7 @@ def handle_int_env_var(var: str, default: int = 0) -> int:
 
 # deployer constants
 DEPLOYER_DOCKER_IMAGE_KEY = "deployer"
+IN_MEMORY_ARTIFACT_URI_PREFIX = "memory://"
 
 # Secret constants
 SECRET_VALUES = "values"

src/zenml/deployers/containerized_deployer.py

@@ -11,7 +11,7 @@
 #  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.
-"""Base class for all ZenML deployers."""
+"""Base class for all containerized deployers."""
 
 from abc import ABC
 from typing import (

src/zenml/deployers/docker/docker_deployer.py

@@ -365,6 +365,7 @@ def do_provision_deployment(
             preferred_ports=preferred_ports,
             allocate_port_if_busy=settings.allocate_port_if_busy,
             range=settings.port_range,
+            address="0.0.0.0",  # nosec
         )
         ports: Dict[str, Optional[int]] = {"8000/tcp": port}
 

src/zenml/deployers/server/runtime.py

@@ -48,10 +48,10 @@ def reset(self) -> None:
         self.active = False
         self.request_id = None
         self.snapshot_id = None
-        self.pipeline_parameters.clear()
-        self.outputs.clear()
+        self.pipeline_parameters = {}
+        self.outputs = {}
         self.skip_artifact_materialization = False
-        self.in_memory_data.clear()
+        self.in_memory_data = {}
 
 
 _deployment_context: contextvars.ContextVar[_DeploymentState] = (
@@ -166,5 +166,5 @@ def get_in_memory_data(uri: str) -> Any:
     """
     if is_active():
         state = _get_context()
-        return state.in_memory_data.get(uri)
+        return state.in_memory_data[uri]
     return None

src/zenml/deployers/server/service.py

@@ -35,6 +35,9 @@
     ServiceInfo,
     SnapshotInfo,
 )
+from zenml.deployers.utils import (
+    deployment_snapshot_request_from_source_snapshot,
+)
 from zenml.enums import StackComponentType
 from zenml.hooks.hook_validators import load_and_run_hook
 from zenml.logger import get_logger
@@ -337,11 +340,6 @@ def _prepare_execute_with_orchestrator(
         Returns:
             A tuple of (placeholder_run, deployment_snapshot).
         """
-        # Create a new snapshot with deployment-specific parameters and settings
-        from zenml.orchestrators.utils import (
-            deployment_snapshot_request_from_source_snapshot,
-        )
-
         deployment_snapshot_request = (
             deployment_snapshot_request_from_source_snapshot(
                 source_snapshot=self.snapshot,

src/zenml/deployers/utils.py

@@ -21,14 +21,21 @@
 import requests
 
 from zenml.client import Client
+from zenml.config.step_configurations import Step
 from zenml.deployers.exceptions import (
     DeploymentHTTPError,
     DeploymentNotFoundError,
     DeploymentProvisionError,
 )
 from zenml.enums import DeploymentStatus
-from zenml.models import DeploymentResponse
+from zenml.models import (
+    CodeReferenceRequest,
+    DeploymentResponse,
+    PipelineSnapshotRequest,
+    PipelineSnapshotResponse,
+)
 from zenml.steps.step_context import get_step_context
+from zenml.utils import pydantic_utils
 from zenml.utils.json_utils import pydantic_encoder
 
 
@@ -274,3 +281,120 @@ def invoke_deployment(
         raise DeploymentHTTPError(
             f"Request failed for deployment {deployment_name_or_id}: {e}"
         )
+
+
+def deployment_snapshot_request_from_source_snapshot(
+    source_snapshot: PipelineSnapshotResponse,
+    deployment_parameters: Dict[str, Any],
+) -> PipelineSnapshotRequest:
+    """Generate a snapshot request for deployment execution.
+
+    Args:
+        source_snapshot: The source snapshot from which to create the
+            snapshot request.
+        deployment_parameters: Parameters to override for deployment execution.
+
+    Raises:
+        RuntimeError: If the source snapshot does not have an associated stack.
+
+    Returns:
+        The generated snapshot request.
+    """
+    if source_snapshot.stack is None:
+        raise RuntimeError("Missing source snapshot stack")
+
+    pipeline_configuration = pydantic_utils.update_model(
+        source_snapshot.pipeline_configuration, {"enable_cache": False}
+    )
+
+    steps = {}
+    for invocation_id, step in source_snapshot.step_configurations.items():
+        updated_step_parameters = step.config.parameters.copy()
+
+        for param_name in step.config.parameters:
+            if param_name in deployment_parameters:
+                updated_step_parameters[param_name] = deployment_parameters[
+                    param_name
+                ]
+
+        # Deployment-specific step overrides
+        step_update = {
+            "enable_cache": False,  # Disable caching for all steps
+            "step_operator": None,  # Remove step operators for deployments
+            "retry": None,  # Remove retry configuration
+            "parameters": updated_step_parameters,
+        }
+
+        step_config = pydantic_utils.update_model(
+            step.step_config_overrides, step_update
+        )
+        merged_step_config = step_config.apply_pipeline_configuration(
+            pipeline_configuration
+        )
+
+        steps[invocation_id] = Step(
+            spec=step.spec,
+            config=merged_step_config,
+            step_config_overrides=step_config,
+        )
+
+    code_reference_request = None
+    if source_snapshot.code_reference:
+        code_reference_request = CodeReferenceRequest(
+            commit=source_snapshot.code_reference.commit,
+            subdirectory=source_snapshot.code_reference.subdirectory,
+            code_repository=source_snapshot.code_reference.code_repository.id,
+        )
+
+    zenml_version = Client().zen_store.get_store_info().version
+
+    # Compute the source snapshot ID:
+    # - If the source snapshot has a name, we use it as the source snapshot.
+    #   That way, all runs will be associated with this snapshot.
+    # - If the source snapshot is based on another snapshot (which therefore
+    #   has a name), we use that one instead.
+    # - If the source snapshot does not have a name and is not based on another
+    #   snapshot, we don't set a source snapshot.
+    #
+    # With this, we ensure that all runs are associated with the closest named
+    # source snapshot.
+    source_snapshot_id = None
+    if source_snapshot.name:
+        source_snapshot_id = source_snapshot.id
+    elif source_snapshot.source_snapshot_id:
+        source_snapshot_id = source_snapshot.source_snapshot_id
+
+    updated_pipeline_spec = source_snapshot.pipeline_spec
+    if (
+        source_snapshot.pipeline_spec
+        and source_snapshot.pipeline_spec.parameters is not None
+    ):
+        original_params: Dict[str, Any] = dict(
+            source_snapshot.pipeline_spec.parameters
+        )
+        merged_params: Dict[str, Any] = original_params.copy()
+        for k, v in deployment_parameters.items():
+            if k in original_params:
+                merged_params[k] = v
+        updated_pipeline_spec = pydantic_utils.update_model(
+            source_snapshot.pipeline_spec, {"parameters": merged_params}
+        )
+
+    return PipelineSnapshotRequest(
+        project=source_snapshot.project_id,
+        run_name_template=source_snapshot.run_name_template,
+        pipeline_configuration=pipeline_configuration,
+        step_configurations=steps,
+        client_environment={},
+        client_version=zenml_version,
+        server_version=zenml_version,
+        stack=source_snapshot.stack.id,
+        pipeline=source_snapshot.pipeline.id,
+        schedule=None,
+        code_reference=code_reference_request,
+        code_path=source_snapshot.code_path,
+        build=source_snapshot.build.id if source_snapshot.build else None,
+        source_snapshot=source_snapshot_id,
+        pipeline_version_hash=source_snapshot.pipeline_version_hash,
+        pipeline_spec=updated_pipeline_spec,
+    )