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

zenml/zen_server/routers/artifact_version_endpoints.py

@@ -13,13 +13,26 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for artifact versions."""
 
+import os
 from typing import List, Union
 from uuid import UUID
 
 from fastapi import APIRouter, Depends, Security
+from fastapi.responses import FileResponse
+from starlette.background import BackgroundTask
 
-from zenml.artifacts.utils import load_artifact_visualization
-from zenml.constants import API, ARTIFACT_VERSIONS, BATCH, VERSION_1, VISUALIZE
+from zenml.artifacts.utils import (
+    load_artifact_visualization,
+)
+from zenml.constants import (
+    API,
+    ARTIFACT_VERSIONS,
+    BATCH,
+    DATA,
+    DOWNLOAD_TOKEN,
+    VERSION_1,
+    VISUALIZE,
+)
 from zenml.models import (
     ArtifactVersionFilter,
     ArtifactVersionRequest,
@@ -28,7 +41,16 @@ from zenml.models import (
     LoadedVisualization,
     Page,
 )
-from zenml.zen_server.auth import AuthContext, authorize
+from zenml.zen_server.auth import (
+    AuthContext,
+    authorize,
+    generate_artifact_download_token,
+    verify_artifact_download_token,
+)
+from zenml.zen_server.download_utils import (
+    create_artifact_archive,
+    verify_artifact_is_downloadable,
+)
 from zenml.zen_server.exceptions import error_response
 from zenml.zen_server.rbac.endpoint_utils import (
     verify_permissions_and_batch_create_entity,
@@ -275,3 +297,64 @@ def get_artifact_visualization(
     return load_artifact_visualization(
         artifact=artifact, index=index, zen_store=store, encode_image=True
     )
+
+
+@artifact_version_router.get(
+    "/{artifact_version_id}" + DOWNLOAD_TOKEN,
+    responses={401: error_response, 404: error_response, 422: error_response},
+)
+@handle_exceptions
+def get_artifact_download_token(
+    artifact_version_id: UUID,
+    _: AuthContext = Security(authorize),
+) -> str:
+    """Get a download token for the artifact data.
+
+    Args:
+        artifact_version_id: ID of the artifact version for which to get the data.
+
+    Returns:
+        The download token for the artifact data.
+    """
+    artifact = verify_permissions_and_get_entity(
+        id=artifact_version_id, get_method=zen_store().get_artifact_version
+    )
+    verify_artifact_is_downloadable(artifact)
+
+    # The artifact download is handled in a separate tab by the browser. In this
+    # tab, we do not have the ability to set any headers and therefore cannot
+    # include the CSRF token in the request. To handle this, we instead generate
+    # a JWT token in this endpoint (which includes CSRF and RBAC checks) and
+    # then use that token to download the artifact data in a separate endpoint
+    # which only verifies this short-lived token.
+    return generate_artifact_download_token(artifact_version_id)
+
+
+@artifact_version_router.get(
+    "/{artifact_version_id}" + DATA,
+    responses={401: error_response, 404: error_response, 422: error_response},
+)
+@handle_exceptions
+def download_artifact_data(
+    artifact_version_id: UUID, token: str
+) -> FileResponse:
+    """Download the artifact data.
+
+    Args:
+        artifact_version_id: ID of the artifact version for which to get the data.
+        token: The token to authenticate the artifact download.
+
+    Returns:
+        The artifact data.
+    """
+    verify_artifact_download_token(token, artifact_version_id)
+
+    artifact = zen_store().get_artifact_version(artifact_version_id)
+    archive_path = create_artifact_archive(artifact)
+
+    return FileResponse(
+        archive_path,
+        media_type="application/gzip",
+        filename=f"{artifact.name}-{artifact.version}.tar.gz",
+        background=BackgroundTask(os.remove, archive_path),
+    )

zenml/zen_server/routers/auth_endpoints.py

@@ -477,7 +477,6 @@ def api_token(
     expires_in: Optional[int] = None,
     schedule_id: Optional[UUID] = None,
     pipeline_run_id: Optional[UUID] = None,
-    step_run_id: Optional[UUID] = None,
     auth_context: AuthContext = Security(authorize),
 ) -> str:
     """Generate an API token for the current user.
@@ -506,7 +505,6 @@ def api_token(
         schedule_id: The ID of the schedule to scope the workload API token to.
         pipeline_run_id: The ID of the pipeline run to scope the workload API
             token to.
-        step_run_id: The ID of the step run to scope the workload API token to.
         auth_context: The authentication context.
 
     Returns:
@@ -522,10 +520,10 @@ def api_token(
         raise AuthorizationException("Not authenticated.")
 
     if token_type == APITokenType.GENERIC:
-        if schedule_id or pipeline_run_id or step_run_id:
+        if schedule_id or pipeline_run_id:
             raise ValueError(
-                "Generic API tokens cannot be scoped to a schedule, pipeline "
-                "run or step run."
+                "Generic API tokens cannot be scoped to a schedule or pipeline "
+                "run."
             )
 
         config = server_config()
@@ -549,12 +547,10 @@ def api_token(
 
     schedule_id = schedule_id or token.schedule_id
     pipeline_run_id = pipeline_run_id or token.pipeline_run_id
-    step_run_id = step_run_id or token.step_run_id
 
-    if not pipeline_run_id and not schedule_id and not step_run_id:
+    if not pipeline_run_id and not schedule_id:
         raise ValueError(
-            "Workload API tokens must be scoped to a schedule, pipeline run "
-            "or step run."
+            "Workload API tokens must be scoped to a schedule or pipeline run."
         )
 
     if schedule_id and token.schedule_id and schedule_id != token.schedule_id:
@@ -575,13 +571,6 @@ def api_token(
             f"pipeline run {token.pipeline_run_id}."
         )
 
-    if step_run_id and token.step_run_id and step_run_id != token.step_run_id:
-        raise AuthorizationException(
-            f"Unable to scope API token to step run {step_run_id}. The "
-            f"token used to authorize this request is already scoped to "
-            f"step run {token.step_run_id}."
-        )
-
     project_id: Optional[UUID] = None
 
     if schedule_id:
@@ -623,25 +612,6 @@ def api_token(
                 "for security reasons."
             )
 
-    if step_run_id:
-        # The step run must exist and the step must not be concluded
-        try:
-            step_run = zen_store().get_run_step(step_run_id, hydrate=False)
-        except KeyError:
-            raise ValueError(
-                f"Step run {step_run_id} does not exist and API tokens cannot "
-                "be generated for non-existent step runs for security reasons."
-            )
-
-        project_id = step_run.project.id
-
-        if step_run.status.is_finished:
-            raise ValueError(
-                f"The execution of step run {step_run_id} has already "
-                "concluded and API tokens can no longer be generated for it "
-                "for security reasons."
-            )
-
     assert project_id is not None
     verify_permission(
         resource_type=ResourceType.PIPELINE_RUN,
@@ -656,7 +626,6 @@ def api_token(
         device=auth_context.device,
         schedule_id=schedule_id,
         pipeline_run_id=pipeline_run_id,
-        step_run_id=step_run_id,
         # Don't include the access token as a cookie in the response
         response=None,
         # Never expire the token

zenml/zen_server/routers/pipeline_deployments_endpoints.py

@@ -19,11 +19,13 @@ from uuid import UUID
 from fastapi import APIRouter, Depends, Security
 
 from zenml.constants import API, PIPELINE_DEPLOYMENTS, VERSION_1
+from zenml.logging.step_logging import fetch_logs
 from zenml.models import (
     Page,
     PipelineDeploymentFilter,
     PipelineDeploymentRequest,
     PipelineDeploymentResponse,
+    PipelineRunFilter,
 )
 from zenml.zen_server.auth import AuthContext, authorize
 from zenml.zen_server.exceptions import error_response
@@ -179,33 +181,68 @@ def delete_deployment(
     )
 
 
-if server_config().workload_manager_enabled:
+@router.get(
+    "/{deployment_id}/logs",
+    responses={
+        401: error_response,
+        404: error_response,
+        422: error_response,
+    },
+)
+@handle_exceptions
+def deployment_logs(
+    deployment_id: UUID,
+    offset: int = 0,
+    length: int = 1024 * 1024 * 16,  # Default to 16MiB of data
+    _: AuthContext = Security(authorize),
+) -> str:
+    """Get deployment logs.
+
+    Args:
+        deployment_id: ID of the deployment.
+        offset: The offset from which to start reading.
+        length: The amount of bytes that should be read.
+
+    Returns:
+        The deployment logs.
+
+    Raises:
+        KeyError: If no logs are available for the deployment.
+    """
+    store = zen_store()
 
-    @router.get(
-        "/{deployment_id}/logs",
-        responses={
-            401: error_response,
-            404: error_response,
-            422: error_response,
-        },
+    deployment = verify_permissions_and_get_entity(
+        id=deployment_id,
+        get_method=store.get_deployment,
+        hydrate=True,
     )
-    @handle_exceptions
-    def deployment_logs(
-        deployment_id: UUID,
-        _: AuthContext = Security(authorize),
-    ) -> str:
-        """Get deployment logs.
-
-        Args:
-            deployment_id: ID of the deployment.
-
-        Returns:
-            The deployment logs.
-        """
-        deployment = verify_permissions_and_get_entity(
-            id=deployment_id,
-            get_method=zen_store().get_deployment,
-            hydrate=True,
-        )
 
+    if deployment.template_id and server_config().workload_manager_enabled:
         return workload_manager().get_logs(workload_id=deployment.id)
+
+    # Get the last pipeline run for this deployment
+    pipeline_runs = store.list_runs(
+        runs_filter_model=PipelineRunFilter(
+            project=deployment.project.id,
+            sort_by="asc:created",
+            size=1,
+            deployment_id=deployment.id,
+        )
+    )
+
+    if len(pipeline_runs.items) == 0:
+        return ""
+
+    run = pipeline_runs.items[0]
+
+    logs = run.logs
+    if logs is None:
+        raise KeyError("No logs available for this deployment")
+
+    return fetch_logs(
+        zen_store=store,
+        artifact_store_id=logs.artifact_store_id,
+        logs_uri=logs.uri,
+        offset=offset,
+        length=length,
+    )

zenml/zen_server/routers/runs_endpoints.py

@@ -29,6 +29,7 @@ from zenml.constants import (
 )
 from zenml.enums import ExecutionStatus, StackComponentType
 from zenml.logger import get_logger
+from zenml.logging.step_logging import fetch_logs
 from zenml.models import (
     Page,
     PipelineRunFilter,
@@ -55,6 +56,8 @@ from zenml.zen_server.routers.projects_endpoints import workspace_router
 from zenml.zen_server.utils import (
     handle_exceptions,
     make_dependable,
+    server_config,
+    workload_manager,
     zen_store,
 )
 
@@ -375,3 +378,57 @@ def refresh_run_status(
             f"The stack, the run '{run.id}' was executed on, is deleted."
         )
     run.refresh_run_status()
+
+
+@router.get(
+    "/{run_id}/logs",
+    responses={
+        401: error_response,
+        404: error_response,
+        422: error_response,
+    },
+)
+@handle_exceptions
+def run_logs(
+    run_id: UUID,
+    offset: int = 0,
+    length: int = 1024 * 1024 * 16,  # Default to 16MiB of data
+    _: AuthContext = Security(authorize),
+) -> str:
+    """Get pipeline run logs.
+
+    Args:
+        run_id: ID of the pipeline run.
+        offset: The offset from which to start reading.
+        length: The amount of bytes that should be read.
+
+    Returns:
+        The pipeline run logs.
+
+    Raises:
+        KeyError: If no logs are available for the pipeline run.
+    """
+    store = zen_store()
+
+    run = verify_permissions_and_get_entity(
+        id=run_id,
+        get_method=store.get_run,
+        hydrate=True,
+    )
+
+    if run.deployment_id:
+        deployment = store.get_deployment(run.deployment_id)
+        if deployment.template_id and server_config().workload_manager_enabled:
+            return workload_manager().get_logs(workload_id=deployment.id)
+
+    logs = run.logs
+    if logs is None:
+        raise KeyError("No logs available for this pipeline run")
+
+    return fetch_logs(
+        zen_store=store,
+        artifact_store_id=logs.artifact_store_id,
+        logs_uri=logs.uri,
+        offset=offset,
+        length=length,
+    )

zenml/zen_server/routers/users_endpoints.py

@@ -698,7 +698,7 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
 if server_config().rbac_enabled:
 
     @router.post(
-        "/{user_name_or_id}/resource_membership",
+        "/resource_membership",
         responses={
             401: error_response,
             404: error_response,
@@ -707,16 +707,16 @@ if server_config().rbac_enabled:
     )
     @handle_exceptions
     def update_user_resource_membership(
-        user_name_or_id: Union[str, UUID],
         resource_type: str,
         resource_id: UUID,
         actions: List[str],
+        user_id: Optional[str] = None,
+        team_id: Optional[str] = None,
         auth_context: AuthContext = Security(authorize),
     ) -> None:
         """Updates resource memberships of a user.
 
         Args:
-            user_name_or_id: Name or ID of the user.
             resource_type: Type of the resource for which to update the
                 membership.
             resource_id: ID of the resource for which to update the membership.
@@ -724,16 +724,19 @@ if server_config().rbac_enabled:
                 the resource. If the user currently has permissions to perform
                 actions which are not passed in this list, the permissions will
                 be removed.
+            user_id: ID of the user for which to update the membership.
+            team_id: ID of the team for which to update the membership.
             auth_context: Authentication context.
 
         Raises:
             ValueError: If a user tries to update their own membership.
             KeyError: If no resource with the given type and ID exists.
         """
-        user = zen_store().get_user(user_name_or_id)
-        # verify_permission_for_model(user, action=Action.READ)
-
-        if user.id == auth_context.user.id:
+        if (
+            user_id
+            and auth_context.user.external_user_id
+            and user_id == str(auth_context.user.external_user_id)
+        ):
             raise ValueError(
                 "Not allowed to call endpoint with the authenticated user."
             )
@@ -765,7 +768,9 @@ if server_config().rbac_enabled:
             verify_permission_for_model(model=model, action=Action(action))
 
         update_resource_membership(
-            user=user,
+            sharing_user=auth_context.user,
             resource=resource,
             actions=[Action(action) for action in actions],
+            user_id=user_id,
+            team_id=team_id,
         )

zenml/zen_server/template_execution/utils.py

@@ -373,7 +373,7 @@ def deployment_request_from_template(
     )
 
     step_config_dict_base = pipeline_configuration.model_dump(
-        exclude={"name", "parameters", "tags"}
+        exclude={"name", "parameters", "tags", "enable_pipeline_logs"}
     )
     steps = {}
     for invocation_id, step in deployment.step_configurations.items():
@@ -411,14 +411,14 @@ def deployment_request_from_template(
         if unknown_parameters:
             raise ValueError(
                 "Run configuration contains the following unknown "
-                f"parameters for step {step.config.name}: {unknown_parameters}."
+                f"parameters for step {invocation_id}: {unknown_parameters}."
             )
 
         missing_parameters = required_parameters - configured_parameters
         if missing_parameters:
             raise ValueError(
                 "Run configuration is missing the following required "
-                f"parameters for step {step.config.name}: {missing_parameters}."
+                f"parameters for step {invocation_id}: {missing_parameters}."
             )
 
         step_config = StepConfiguration.model_validate(step_config_dict)

zenml/zen_stores/migrations/versions/ff538a321a92_migrate_onboarding_state.py

@@ -0,0 +1,123 @@
+"""Migrate onboarding state [ff538a321a92].
+
+Revision ID: ff538a321a92
+Revises: 0.80.2
+Create Date: 2025-04-11 09:30:03.324310
+
+"""
+
+import json
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "ff538a321a92"
+down_revision = "0.80.2"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    with op.batch_alter_table("server_settings", schema=None) as batch_op:
+        batch_op.alter_column(
+            "onboarding_state",
+            existing_type=sa.VARCHAR(),
+            type_=sa.TEXT(),
+            existing_nullable=True,
+        )
+
+    connection = op.get_bind()
+
+    meta = sa.MetaData()
+    meta.reflect(only=("server_settings",), bind=connection)
+
+    server_settings_table = sa.Table("server_settings", meta)
+
+    existing_onboarding_state = connection.execute(
+        sa.select(server_settings_table.c.onboarding_state)
+    ).scalar_one_or_none()
+
+    if not existing_onboarding_state:
+        return
+
+    state = json.loads(existing_onboarding_state)
+
+    meta = sa.MetaData()
+    meta.reflect(
+        only=(
+            "pipeline_run",
+            "stack_component",
+            "stack",
+            "stack_composition",
+            "pipeline_deployment",
+        ),
+        bind=connection,
+    )
+
+    pipeline_run_table = sa.Table("pipeline_run", meta)
+    stack_component_table = sa.Table("stack_component", meta)
+    stack_table = sa.Table("stack", meta)
+    stack_composition_table = sa.Table("stack_composition", meta)
+    pipeline_deployment_table = sa.Table("pipeline_deployment", meta)
+
+    stack_with_remote_artifact_store_count = connection.execute(
+        sa.select(sa.func.count(stack_table.c.id))
+        .where(stack_composition_table.c.stack_id == stack_table.c.id)
+        .where(
+            stack_composition_table.c.component_id
+            == stack_component_table.c.id
+        )
+        .where(stack_component_table.c.flavor != "local")
+        .where(stack_component_table.c.type == "artifact_store")
+    ).scalar()
+    if (
+        stack_with_remote_artifact_store_count
+        and stack_with_remote_artifact_store_count > 0
+    ):
+        state.append("stack_with_remote_artifact_store_created")
+
+    pipeline_run_with_remote_artifact_store_count = connection.execute(
+        sa.select(sa.func.count(pipeline_run_table.c.id))
+        .where(
+            pipeline_run_table.c.deployment_id
+            == pipeline_deployment_table.c.id
+        )
+        .where(pipeline_deployment_table.c.stack_id == stack_table.c.id)
+        .where(stack_composition_table.c.stack_id == stack_table.c.id)
+        .where(
+            stack_composition_table.c.component_id
+            == stack_component_table.c.id
+        )
+        .where(stack_component_table.c.flavor != "local")
+        .where(stack_component_table.c.type == "artifact_store")
+    ).scalar()
+    if (
+        pipeline_run_with_remote_artifact_store_count
+        and pipeline_run_with_remote_artifact_store_count > 0
+    ):
+        state.append("pipeline_run_with_remote_artifact_store")
+        state.append("production_setup_completed")
+
+    # Remove duplicate keys
+    state = list(set(state))
+
+    connection.execute(
+        sa.update(server_settings_table).values(
+            onboarding_state=json.dumps(state)
+        )
+    )
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("server_settings", schema=None) as batch_op:
+        batch_op.alter_column(
+            "onboarding_state",
+            existing_type=sa.TEXT(),
+            type_=sa.VARCHAR(),
+            existing_nullable=True,
+        )
+    # ### end Alembic commands ###

zenml/zen_stores/rest_zen_store.py

@@ -4074,15 +4074,6 @@ class RestZenStore(BaseZenStore):
             # Check if username and password are configured
             username, password = credentials_store.get_password(self.url)
 
-            api_key_hint = (
-                "\nHint: If you're getting this error in an automated, "
-                "non-interactive workload like a pipeline run or a CI/CD job, "
-                "you should use a service account API key to authenticate to "
-                "the server instead of temporary CLI login credentials. For "
-                "more information, see "
-                "https://docs.zenml.io/how-to/project-setup-and-management/connecting-to-zenml/connect-with-a-service-account"
-            )
-
             if api_key is not None:
                 # An API key is configured. Use it as a password to
                 # authenticate.
@@ -4119,14 +4110,12 @@ class RestZenStore(BaseZenStore):
                         "You need to be logged in to ZenML Pro in order to "
                         f"access the ZenML Pro server '{self.url}'. Please run "
                         "'zenml login' to log in or choose a different server."
-                        + api_key_hint
                     )
 
                 elif pro_token.expired:
                     raise CredentialsNotValid(
                         "Your ZenML Pro login session has expired. "
                         "Please log in again using 'zenml login'."
-                        + api_key_hint
                     )
 
                 data = {
@@ -4140,13 +4129,12 @@ class RestZenStore(BaseZenStore):
                     raise CredentialsNotValid(
                         "No valid credentials found. Please run 'zenml login "
                         f"--url {self.url}' to connect to the current server."
-                        + api_key_hint
                     )
                 elif token.expired:
                     raise CredentialsNotValid(
                         "Your authentication to the current server has expired. "
                         "Please log in again using 'zenml login --url "
-                        f"{self.url}'." + api_key_hint
+                        f"{self.url}'."
                     )
 
             response = self._handle_response(
@@ -4405,6 +4393,7 @@ class RestZenStore(BaseZenStore):
                 # explicitly indicates that the credentials are not valid and
                 # they can be thrown away or when the request is not
                 # authenticated at all.
+                credentials_store = get_credentials_store()
 
                 if self._api_token is None:
                     # The last request was not authenticated with an API
@@ -4416,6 +4405,20 @@ class RestZenStore(BaseZenStore):
                         "Re-authenticating and retrying..."
                     )
                     self.authenticate()
+                elif not credentials_store.can_login(self.url):
+                    # The request failed either because we're not
+                    # authenticated or our current credentials are not valid
+                    # anymore.
+                    logger.error(
+                        "The current token is no longer valid, and "
+                        "it is not possible to generate a new token using the "
+                        "configured credentials. Please run "
+                        f"`zenml login --url {self.url}` to re-authenticate to "
+                        "the server or authenticate using an API key. See "
+                        "https://docs.zenml.io/how-to/project-setup-and-management/connecting-to-zenml/connect-with-a-service-account "
+                        "for more information."
+                    )
+                    raise e
                 elif not re_authenticated:
                     # The last request was authenticated with an API token
                     # that was rejected by the server. We attempt a

zenml/zen_stores/schemas/pipeline_run_schemas.py

@@ -412,6 +412,7 @@ class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True):
             resources = PipelineRunResponseResources(
                 model_version=model_version,
                 tags=[tag.to_model() for tag in self.tags],
+                logs=self.logs.to_model() if self.logs else None,
             )
 
         return PipelineRunResponse(