zenml-nightly 0.80.2.dev20250414__py3-none-any.whl → 0.80.2.dev20250416__py3-none-any.whl

zenml/pipelines/run_utils.py

@@ -15,6 +15,7 @@ from zenml.enums import ExecutionStatus
 from zenml.logger import get_logger
 from zenml.models import (
     FlavorFilter,
+    LogsRequest,
     PipelineDeploymentBase,
     PipelineDeploymentResponse,
     PipelineRunRequest,
@@ -49,6 +50,7 @@ def get_default_run_name(pipeline_name: str) -> str:
 
 def create_placeholder_run(
     deployment: "PipelineDeploymentResponse",
+    logs: Optional["LogsRequest"] = None,
 ) -> Optional["PipelineRunResponse"]:
     """Create a placeholder run for the deployment.
 
@@ -57,6 +59,7 @@ def create_placeholder_run(
 
     Args:
         deployment: The deployment for which to create the placeholder run.
+        logs: The logs for the run.
 
     Returns:
         The placeholder run or `None` if no run was created.
@@ -86,6 +89,7 @@ def create_placeholder_run(
         pipeline=deployment.pipeline.id if deployment.pipeline else None,
         status=ExecutionStatus.INITIALIZING,
         tags=deployment.pipeline_configuration.tags,
+        logs=logs,
     )
     run, _ = Client().zen_store.get_or_create_run(run_request)
     return run

zenml/steps/utils.py

@@ -553,7 +553,7 @@ def run_as_single_step_pipeline(
     orchestrator = Client().active_stack.orchestrator
 
     pipeline_settings: Any = {}
-    if "synchronous" in orchestrator.config.model_fields:
+    if "synchronous" in type(orchestrator.config).model_fields:
         # Make sure the orchestrator runs sync so we stream the logs
         key = settings_utils.get_stack_component_setting_key(orchestrator)
         pipeline_settings[key] = BaseSettings(synchronous=True)

zenml/utils/io_utils.py

@@ -205,6 +205,29 @@ def resolve_relative_path(path: str) -> str:
     return str(Path(path).resolve())
 
 
+def is_path_within_directory(path: str, directory: str) -> bool:
+    """Checks if a path is contained within a given directory.
+
+    This utility function verifies that a path (absolute or relative) resolves
+    to a location that is within the specified directory. This is useful for
+    security checks such as preventing path traversal attacks when extracting
+    archives (CVE-2007-4559) or whenever path containment needs to be verified.
+
+    Args:
+        path: The path to check (can be relative or absolute).
+        directory: The directory that should contain the path.
+
+    Returns:
+        Boolean indicating whether the path is contained within the directory (True)
+        or not (False).
+    """
+    # Convert to absolute path, ensuring it's normalized
+    abs_path = os.path.abspath(os.path.join(directory, path))
+    # Check if the path is within the target directory
+    dir_abs = os.path.abspath(directory)
+    return abs_path.startswith(dir_abs + os.sep) or abs_path == dir_abs
+
+
 def move(source: str, destination: str, overwrite: bool = False) -> None:
     """Moves dir or file from source to destination. Can be used to rename.
 

zenml/zen_server/auth.py

@@ -15,7 +15,7 @@
 
 from contextvars import ContextVar
 from datetime import datetime, timedelta
-from typing import Callable, Optional, Union
+from typing import Callable, Optional, Tuple, Union
 from urllib.parse import urlencode, urlparse
 from uuid import UUID, uuid4
 
@@ -33,9 +33,12 @@ from zenml.analytics.context import AnalyticsContext
 from zenml.constants import (
     API,
     DEFAULT_USERNAME,
+    DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_LIFETIME,
+    ENV_ZENML_WORKLOAD_TOKEN_EXPIRATION_LEEWAY,
     EXTERNAL_AUTHENTICATOR_TIMEOUT,
     LOGIN,
     VERSION_1,
+    handle_int_env_var,
 )
 from zenml.enums import (
     AuthScheme,
@@ -420,28 +423,32 @@ def authenticate_credentials(
             @cache_result(expiry=30)
             def get_pipeline_run_status(
                 pipeline_run_id: UUID,
-            ) -> Optional[ExecutionStatus]:
+            ) -> Tuple[Optional[ExecutionStatus], Optional[datetime]]:
                 """Get the status of a pipeline run.
 
                 Args:
                     pipeline_run_id: The pipeline run ID.
 
                 Returns:
-                    The pipeline run status or None if the pipeline run does not
-                    exist.
+                    The pipeline run status and end time or None if the pipeline
+                    run does not exist.
                 """
                 try:
                     pipeline_run = zen_store().get_run(
-                        pipeline_run_id, hydrate=False
+                        pipeline_run_id, hydrate=True
                     )
                 except KeyError:
-                    return None
+                    return None, None
 
-                return pipeline_run.status
+                return (
+                    pipeline_run.status,
+                    pipeline_run.end_time,
+                )
 
-            pipeline_run_status = get_pipeline_run_status(
-                decoded_token.pipeline_run_id
-            )
+            (
+                pipeline_run_status,
+                pipeline_run_end_time,
+            ) = get_pipeline_run_status(decoded_token.pipeline_run_id)
             if pipeline_run_status is None:
                 error = (
                     f"Authentication error: error retrieving token pipeline run "
@@ -450,59 +457,35 @@ def authenticate_credentials(
                 logger.error(error)
                 raise CredentialsNotValid(error)
 
+            leeway = handle_int_env_var(
+                ENV_ZENML_WORKLOAD_TOKEN_EXPIRATION_LEEWAY,
+                DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_LIFETIME,
+            )
             if pipeline_run_status.is_finished:
-                error = (
-                    f"The execution of pipeline run "
-                    f"{decoded_token.pipeline_run_id} has already concluded and "
-                    "API tokens scoped to it are no longer valid."
-                )
-                logger.error(error)
-                raise CredentialsNotValid(error)
-
-        if decoded_token.step_run_id:
-            # If the token contains a step run ID, we need to check if the
-            # step run exists in the database and the step run has not concluded.
-            # We use a cached version of the step run status to avoid unnecessary
-            # database queries.
-
-            @cache_result(expiry=30)
-            def get_step_run_status(
-                step_run_id: UUID,
-            ) -> Optional[ExecutionStatus]:
-                """Get the status of a step run.
-
-                Args:
-                    step_run_id: The step run ID.
-
-                Returns:
-                    The step run status or None if the step run does not exist.
-                """
-                try:
-                    step_run = zen_store().get_run_step(
-                        step_run_id, hydrate=False
+                if leeway < 0:
+                    # The token should never expire, we don't need to check
+                    # the end time.
+                    pass
+                elif (
+                    # We don't know the end time. This should never happen, but
+                    # just in case we always expire the token.
+                    pipeline_run_end_time is None
+                    # Calculate whether the token has expired.
+                    or utc_now(tz_aware=pipeline_run_end_time)
+                    > pipeline_run_end_time + timedelta(seconds=leeway)
+                ):
+                    error = (
+                        f"The pipeline run {decoded_token.pipeline_run_id} has "
+                        "finished and API tokens scoped to it are no longer "
+                        "valid. If you want to increase the expiration time "
+                        "of the token to allow steps to continue for longer "
+                        "after other steps have failed, you can do so by "
+                        "configuring the "
+                        f"`{ENV_ZENML_WORKLOAD_TOKEN_EXPIRATION_LEEWAY}` "
+                        "ZenML server environment variable."
                     )
-                except KeyError:
-                    return None
-
-                return step_run.status
-
-            step_run_status = get_step_run_status(decoded_token.step_run_id)
-            if step_run_status is None:
-                error = (
-                    f"Authentication error: error retrieving token step run "
-                    f"{decoded_token.step_run_id}"
-                )
-                logger.error(error)
-                raise CredentialsNotValid(error)
-
-            if step_run_status.is_finished:
-                error = (
-                    f"The execution of step run "
-                    f"{decoded_token.step_run_id} has already concluded and "
-                    "API tokens scoped to it are no longer valid."
-                )
-                logger.error(error)
-                raise CredentialsNotValid(error)
+                    logger.error(error)
+                    raise CredentialsNotValid(error)
 
         auth_context = AuthContext(
             user=user_model,
@@ -861,7 +844,6 @@ def generate_access_token(
     expires_in: Optional[int] = None,
     schedule_id: Optional[UUID] = None,
     pipeline_run_id: Optional[UUID] = None,
-    step_run_id: Optional[UUID] = None,
 ) -> OAuthTokenResponse:
     """Generates an access token for the given user.
 
@@ -880,7 +862,6 @@ def generate_access_token(
             expire.
         schedule_id: The ID of the schedule to scope the token to.
         pipeline_run_id: The ID of the pipeline run to scope the token to.
-        step_run_id: The ID of the step run to scope the token to.
 
     Returns:
         An authentication response with an access token.
@@ -956,7 +937,6 @@ def generate_access_token(
         api_key_id=api_key.id if api_key else None,
         schedule_id=schedule_id,
         pipeline_run_id=pipeline_run_id,
-        step_run_id=step_run_id,
         # Set the session ID if this is a cross-site request
         session_id=session_id,
     ).encode(expires=expires)
@@ -983,6 +963,58 @@ def generate_access_token(
     )
 
 
+def generate_artifact_download_token(artifact_version_id: UUID) -> str:
+    """Generate a JWT token for artifact download.
+
+    Args:
+        artifact_version_id: The ID of the artifact version to download.
+
+    Returns:
+        The JWT token for the artifact download.
+    """
+    import jwt
+
+    config = server_config()
+
+    return jwt.encode(
+        {
+            "exp": utc_now() + timedelta(seconds=30),
+            "artifact_version_id": str(artifact_version_id),
+        },
+        key=config.jwt_secret_key,
+        algorithm=config.jwt_token_algorithm,
+    )
+
+
+def verify_artifact_download_token(
+    token: str, artifact_version_id: UUID
+) -> None:
+    """Verify a JWT token for artifact download.
+
+    Args:
+        token: The JWT token to verify.
+        artifact_version_id: The ID of the artifact version to download.
+
+    Raises:
+        CredentialsNotValid: If the token is invalid or the artifact version
+            ID does not match.
+    """
+    import jwt
+
+    config = server_config()
+    try:
+        claims = jwt.decode(
+            token,
+            config.jwt_secret_key,
+            algorithms=[config.jwt_token_algorithm],
+        )
+    except jwt.PyJWTError as e:
+        raise CredentialsNotValid(f"Invalid JWT token: {e}") from e
+
+    if claims["artifact_version_id"] != str(artifact_version_id):
+        raise CredentialsNotValid("Invalid artifact version ID")
+
+
 def http_authentication(
     credentials: HTTPBasicCredentials = Depends(HTTPBasic()),
 ) -> AuthContext:

zenml/zen_server/cloud_utils.py

@@ -7,7 +7,10 @@ import requests
 from requests.adapters import HTTPAdapter, Retry
 
 from zenml.config.server_config import ServerProConfiguration
-from zenml.exceptions import SubscriptionUpgradeRequiredError
+from zenml.exceptions import (
+    IllegalOperationError,
+    SubscriptionUpgradeRequiredError,
+)
 from zenml.utils.time_utils import utc_now
 from zenml.zen_server.utils import get_zenml_headers, server_config
 
@@ -43,6 +46,7 @@ class ZenMLCloudConnection:
         Raises:
             SubscriptionUpgradeRequiredError: If the current subscription tier
                 is insufficient for the attempted operation.
+            IllegalOperationError: If the request failed with a 403 status code.
             RuntimeError: If the request failed.
 
         Returns:
@@ -65,6 +69,8 @@ class ZenMLCloudConnection:
         except requests.HTTPError as e:
             if response.status_code == 402:
                 raise SubscriptionUpgradeRequiredError(response.json())
+            elif response.status_code == 403:
+                raise IllegalOperationError(response.json())
             else:
                 raise RuntimeError(
                     f"Failed while trying to contact the central zenml pro "

zenml/zen_server/download_utils.py

@@ -0,0 +1,123 @@
+#  Copyright (c) ZenML GmbH 2025. 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.
+"""Utility functions for downloading artifacts."""
+
+import os
+import tarfile
+import tempfile
+from typing import (
+    TYPE_CHECKING,
+)
+
+from zenml.artifacts.utils import _load_artifact_store
+from zenml.exceptions import (
+    IllegalOperationError,
+)
+from zenml.models import (
+    ArtifactVersionResponse,
+)
+from zenml.zen_server.utils import server_config, zen_store
+
+if TYPE_CHECKING:
+    from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
+
+
+def verify_artifact_is_downloadable(
+    artifact: "ArtifactVersionResponse",
+) -> "BaseArtifactStore":
+    """Verify that the given artifact is downloadable.
+
+    Args:
+        artifact: The artifact to verify.
+
+    Raises:
+        IllegalOperationError: If the artifact is too large to be archived.
+        KeyError: If the artifact store is not found or the artifact URI does
+            not exist.
+
+    Returns:
+        The artifact store.
+    """
+    if not artifact.artifact_store_id:
+        raise KeyError(
+            f"Artifact '{artifact.id}' cannot be downloaded because the "
+            "underlying artifact store was deleted."
+        )
+
+    artifact_store = _load_artifact_store(
+        artifact_store_id=artifact.artifact_store_id, zen_store=zen_store()
+    )
+
+    if not artifact_store.exists(artifact.uri):
+        raise KeyError(f"The artifact URI '{artifact.uri}' does not exist.")
+
+    size = artifact_store.size(artifact.uri)
+    max_download_size = server_config().file_download_size_limit
+
+    if size and size > max_download_size:
+        raise IllegalOperationError(
+            f"The artifact '{artifact.id}' is too large to be downloaded. "
+            f"The maximum download size allowed by your ZenML server is "
+            f"{max_download_size} bytes."
+        )
+
+    return artifact_store
+
+
+def create_artifact_archive(
+    artifact: "ArtifactVersionResponse",
+) -> str:
+    """Create an archive of the given artifact.
+
+    Args:
+        artifact: The artifact to archive.
+
+    Returns:
+        The path to the created archive.
+    """
+    artifact_store = verify_artifact_is_downloadable(artifact)
+
+    def _prepare_tarinfo(path: str) -> tarfile.TarInfo:
+        archive_path = os.path.relpath(path, artifact.uri)
+        tarinfo = tarfile.TarInfo(name=archive_path)
+        if size := artifact_store.size(path):
+            tarinfo.size = size
+        return tarinfo
+
+    with tempfile.NamedTemporaryFile(delete=False) as temp_file:
+        with tarfile.open(fileobj=temp_file, mode="w:gz") as tar:
+            if artifact_store.isdir(artifact.uri):
+                for dir, _, files in artifact_store.walk(artifact.uri):
+                    dir = dir.decode() if isinstance(dir, bytes) else dir
+                    dir_info = tarfile.TarInfo(
+                        name=os.path.relpath(dir, artifact.uri)
+                    )
+                    dir_info.type = tarfile.DIRTYPE
+                    dir_info.mode = 0o755
+                    tar.addfile(dir_info)
+
+                    for file in files:
+                        file = (
+                            file.decode() if isinstance(file, bytes) else file
+                        )
+                        path = os.path.join(dir, file)
+                        tarinfo = _prepare_tarinfo(path)
+                        with artifact_store.open(path, "rb") as f:
+                            tar.addfile(tarinfo, fileobj=f)
+            else:
+                tarinfo = _prepare_tarinfo(artifact.uri)
+                with artifact_store.open(artifact.uri, "rb") as f:
+                    tar.addfile(tarinfo, fileobj=f)
+
+        return temp_file.name

zenml/zen_server/jwt.py

@@ -54,7 +54,6 @@ class JWTToken(BaseModel):
     api_key_id: Optional[UUID] = None
     schedule_id: Optional[UUID] = None
     pipeline_run_id: Optional[UUID] = None
-    step_run_id: Optional[UUID] = None
     session_id: Optional[UUID] = None
     claims: Dict[str, Any] = {}
 
@@ -148,16 +147,6 @@ class JWTToken(BaseModel):
                     "UUID"
                 )
 
-        step_run_id: Optional[UUID] = None
-        if "step_run_id" in claims:
-            try:
-                step_run_id = UUID(claims.pop("step_run_id"))
-            except ValueError:
-                raise CredentialsNotValid(
-                    "Invalid JWT token: the step_run_id claim is not a valid "
-                    "UUID"
-                )
-
         session_id: Optional[UUID] = None
         if "session_id" in claims:
             try:
@@ -174,7 +163,6 @@ class JWTToken(BaseModel):
             api_key_id=api_key_id,
             schedule_id=schedule_id,
             pipeline_run_id=pipeline_run_id,
-            step_run_id=step_run_id,
             session_id=session_id,
             claims=claims,
         )
@@ -212,8 +200,6 @@ class JWTToken(BaseModel):
             claims["schedule_id"] = str(self.schedule_id)
         if self.pipeline_run_id:
             claims["pipeline_run_id"] = str(self.pipeline_run_id)
-        if self.step_run_id:
-            claims["step_run_id"] = str(self.step_run_id)
         if self.session_id:
             claims["session_id"] = str(self.session_id)
 

zenml/zen_server/rbac/rbac_interface.py

@@ -14,7 +14,7 @@
 """RBAC interface definition."""
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Dict, List, Set, Tuple
+from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple
 
 from zenml.zen_server.rbac.models import Action, Resource
 
@@ -63,15 +63,22 @@ class RBACInterface(ABC):
 
     @abstractmethod
     def update_resource_membership(
-        self, user: "UserResponse", resource: Resource, actions: List[Action]
+        self,
+        sharing_user: "UserResponse",
+        resource: Resource,
+        actions: List[Action],
+        user_id: Optional[str] = None,
+        team_id: Optional[str] = None,
     ) -> None:
         """Update the resource membership of a user.
 
         Args:
-            user: User for which the resource membership should be updated.
+            sharing_user: User that is sharing the resource.
             resource: The resource.
             actions: The actions that the user should be able to perform on the
                 resource.
+            user_id: ID of the user for which to update the membership.
+            team_id: ID of the team for which to update the membership.
         """
 
     @abstractmethod

zenml/zen_server/rbac/utils.py

@@ -688,21 +688,31 @@ def get_schema_for_resource_type(
 
 
 def update_resource_membership(
-    user: UserResponse, resource: Resource, actions: List[Action]
+    sharing_user: "UserResponse",
+    resource: Resource,
+    actions: List[Action],
+    user_id: Optional[str] = None,
+    team_id: Optional[str] = None,
 ) -> None:
     """Update the resource membership of a user.
 
     Args:
-        user: User for which the resource membership should be updated.
+        sharing_user: User that is sharing the resource.
         resource: The resource.
         actions: The actions that the user should be able to perform on the
             resource.
+        user_id: ID of the user for which to update the membership.
+        team_id: ID of the team for which to update the membership.
     """
     if not server_config().rbac_enabled:
         return
 
     rbac().update_resource_membership(
-        user=user, resource=resource, actions=actions
+        sharing_user=sharing_user,
+        resource=resource,
+        actions=actions,
+        user_id=user_id,
+        team_id=team_id,
     )
 
 

zenml/zen_server/rbac/zenml_cloud_rbac.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Cloud RBAC implementation."""
 
-from typing import TYPE_CHECKING, Dict, List, Set, Tuple
+from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple
 
 from zenml.zen_server.cloud_utils import cloud_connection
 from zenml.zen_server.rbac.models import Action, Resource
@@ -117,22 +117,28 @@ class ZenMLCloudRBAC(RBACInterface):
         return full_resource_access, allowed_ids
 
     def update_resource_membership(
-        self, user: "UserResponse", resource: Resource, actions: List[Action]
+        self,
+        sharing_user: "UserResponse",
+        resource: Resource,
+        actions: List[Action],
+        user_id: Optional[str] = None,
+        team_id: Optional[str] = None,
     ) -> None:
         """Update the resource membership of a user.
 
         Args:
-            user: User for which the resource membership should be updated.
+            sharing_user: User that is sharing the resource.
             resource: The resource.
             actions: The actions that the user should be able to perform on the
                 resource.
+            user_id: ID of the user for which to update the membership.
+            team_id: ID of the team for which to update the membership.
         """
-        if user.is_service_account:
-            # Service accounts have full permissions for now
-            return
-
+        assert sharing_user.external_user_id
         data = {
-            "user_id": str(user.external_user_id),
+            "user_id": user_id,
+            "team_id": team_id,
+            "sharing_user_id": str(sharing_user.external_user_id),
             "resource": str(resource),
             "actions": [str(action) for action in actions],
         }