zenml-nightly 0.68.1.dev20241106__py3-none-any.whl → 0.70.0.dev20241116__py3-none-any.whl

zenml/zen_server/routers/auth_endpoints.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for authentication (login)."""
 
-from datetime import datetime, timedelta
 from typing import Optional, Union
 from urllib.parse import urlencode
 from uuid import UUID
@@ -40,17 +39,17 @@ from zenml.constants import (
     VERSION_1,
 )
 from zenml.enums import (
+    APITokenType,
     AuthScheme,
+    ExecutionStatus,
     OAuthDeviceStatus,
     OAuthGrantTypes,
 )
 from zenml.exceptions import AuthorizationException
 from zenml.logger import get_logger
 from zenml.models import (
-    APIKeyInternalResponse,
     OAuthDeviceAuthorizationResponse,
     OAuthDeviceInternalRequest,
-    OAuthDeviceInternalResponse,
     OAuthDeviceInternalUpdate,
     OAuthDeviceUserAgentHeader,
     OAuthRedirectResponse,
@@ -63,9 +62,9 @@ from zenml.zen_server.auth import (
     authenticate_device,
     authenticate_external_user,
     authorize,
+    generate_access_token,
 )
 from zenml.zen_server.exceptions import error_response
-from zenml.zen_server.jwt import JWTToken
 from zenml.zen_server.rate_limit import rate_limit_requests
 from zenml.zen_server.rbac.models import Action, ResourceType
 from zenml.zen_server.rbac.utils import verify_permission
@@ -211,68 +210,6 @@ class OAuthLoginRequestForm:
                 )
 
 
-def generate_access_token(
-    user_id: UUID,
-    response: Response,
-    device: Optional[OAuthDeviceInternalResponse] = None,
-    api_key: Optional[APIKeyInternalResponse] = None,
-) -> OAuthTokenResponse:
-    """Generates an access token for the given user.
-
-    Args:
-        user_id: The ID of the user.
-        response: The FastAPI response object.
-        device: The device used for authentication.
-        api_key: The service account API key used for authentication.
-
-    Returns:
-        An authentication response with an access token.
-    """
-    config = server_config()
-
-    # The JWT tokens are set to expire according to the values configured
-    # in the server config. Device tokens are handled separately from regular
-    # user tokens.
-    expires: Optional[datetime] = None
-    expires_in: Optional[int] = None
-    if device:
-        # If a device was used for authentication, the token will expire
-        # at the same time as the device.
-        expires = device.expires
-        if expires:
-            expires_in = max(
-                int(expires.timestamp() - datetime.utcnow().timestamp()), 0
-            )
-    elif config.jwt_token_expire_minutes:
-        expires = datetime.utcnow() + timedelta(
-            minutes=config.jwt_token_expire_minutes
-        )
-        expires_in = config.jwt_token_expire_minutes * 60
-
-    access_token = JWTToken(
-        user_id=user_id,
-        device_id=device.id if device else None,
-        api_key_id=api_key.id if api_key else None,
-    ).encode(expires=expires)
-
-    if not device:
-        # Also set the access token as an HTTP only cookie in the response
-        response.set_cookie(
-            key=config.get_auth_cookie_name(),
-            value=access_token,
-            httponly=True,
-            samesite="lax",
-            max_age=config.jwt_token_expire_minutes * 60
-            if config.jwt_token_expire_minutes
-            else None,
-            domain=config.auth_cookie_domain,
-        )
-
-    return OAuthTokenResponse(
-        access_token=access_token, expires_in=expires_in, token_type="bearer"
-    )
-
-
 @router.post(
     LOGIN,
     response_model=Union[OAuthTokenResponse, OAuthRedirectResponse],
@@ -514,45 +451,74 @@ def device_authorization(
 )
 @handle_exceptions
 def api_token(
-    pipeline_id: Optional[UUID] = None,
+    token_type: APITokenType = APITokenType.GENERIC,
     schedule_id: Optional[UUID] = None,
-    expires_minutes: Optional[int] = None,
+    pipeline_run_id: Optional[UUID] = None,
+    step_run_id: Optional[UUID] = None,
     auth_context: AuthContext = Security(authorize),
 ) -> str:
-    """Get a workload API token for the current user.
+    """Generate an API token for the current user.
+
+    Use this endpoint to generate an API token for the current user. Two types
+    of API tokens are supported:
+
+    * Generic API token: This token is short-lived and can be used for
+    generic automation tasks.
+    * Workload API token: This token is scoped to a specific pipeline run, step
+    run or schedule and is used by pipeline workloads to authenticate with the
+    server. A pipeline run ID, step run ID or schedule ID must be provided and
+    the generated token will only be valid for the indicated pipeline run, step
+    run or schedule. No time limit is imposed on the validity of the token.
+    A workload API token can be used to authenticate and generate another
+    workload API token, but only for the same schedule, pipeline run ID or step
+    run ID, in that order.
 
     Args:
-        pipeline_id: The ID of the pipeline to get the API token for.
-        schedule_id: The ID of the schedule to get the API token for.
-        expires_minutes: The number of minutes for which the API token should
-            be valid. If not provided, the API token will be valid indefinitely.
+        token_type: The type of API token to generate.
+        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:
         The API token.
 
     Raises:
-        HTTPException: If the user is not authenticated.
-        AuthorizationException: If trying to scope the API token to a different
-            pipeline/schedule than the token used to authorize this request.
+        AuthorizationException: If not authorized to generate the API token.
+        ValueError: If the request is invalid.
     """
     token = auth_context.access_token
     if not token or not auth_context.encoded_access_token:
         # Should not happen
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Not authenticated.",
-        )
+        raise AuthorizationException("Not authenticated.")
+
+    if token_type == APITokenType.GENERIC:
+        if schedule_id or pipeline_run_id or step_run_id:
+            raise ValueError(
+                "Generic API tokens cannot be scoped to a schedule, pipeline "
+                "run or step run."
+            )
+
+        config = server_config()
+
+        return generate_access_token(
+            user_id=token.user_id,
+            expires_in=config.generic_api_token_lifetime,
+        ).access_token
 
     verify_permission(
         resource_type=ResourceType.PIPELINE_RUN, action=Action.CREATE
     )
 
-    if pipeline_id and token.pipeline_id and pipeline_id != token.pipeline_id:
-        raise AuthorizationException(
-            f"Unable to scope API token to pipeline {pipeline_id}. The "
-            f"token used to authorize this request is already scoped to "
-            f"pipeline {token.pipeline_id}."
+    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:
+        raise ValueError(
+            "Workload API tokens must be scoped to a schedule, pipeline run "
+            "or step run."
         )
 
     if schedule_id and token.schedule_id and schedule_id != token.schedule_id:
@@ -562,21 +528,87 @@ def api_token(
             f"schedule {token.schedule_id}."
         )
 
-    if not token.device_id and not token.api_key_id:
-        # If not authenticated with a device or a service account, the current
-        # API token is returned as is, without any modifications. Issuing
-        # workload tokens is only supported for device authenticated users and
-        # service accounts, because device tokens can be revoked at any time and
-        # service accounts can be disabled.
-        return auth_context.encoded_access_token
-
-    # If authenticated with a device, a new API token is generated for the
-    # pipeline and/or schedule.
-    if pipeline_id:
-        token.pipeline_id = pipeline_id
+    if (
+        pipeline_run_id
+        and token.pipeline_run_id
+        and pipeline_run_id != token.pipeline_run_id
+    ):
+        raise AuthorizationException(
+            f"Unable to scope API token to pipeline run {pipeline_run_id}. The "
+            f"token used to authorize this request is already scoped to "
+            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}."
+        )
+
     if schedule_id:
-        token.schedule_id = schedule_id
-    expires: Optional[datetime] = None
-    if expires_minutes:
-        expires = datetime.utcnow() + timedelta(minutes=expires_minutes)
-    return token.encode(expires=expires)
+        # The schedule must exist
+        try:
+            schedule = zen_store().get_schedule(schedule_id, hydrate=False)
+        except KeyError:
+            raise ValueError(
+                f"Schedule {schedule_id} does not exist and API tokens cannot "
+                "be generated for non-existent schedules for security reasons."
+            )
+
+        if not schedule.active:
+            raise ValueError(
+                f"Schedule {schedule_id} is not active and API tokens cannot "
+                "be generated for inactive schedules for security reasons."
+            )
+
+    if pipeline_run_id:
+        # The pipeline run must exist and the run must not be concluded
+        try:
+            pipeline_run = zen_store().get_run(pipeline_run_id, hydrate=False)
+        except KeyError:
+            raise ValueError(
+                f"Pipeline run {pipeline_run_id} does not exist and API tokens "
+                "cannot be generated for non-existent pipeline runs for "
+                "security reasons."
+            )
+
+        if pipeline_run.status in [
+            ExecutionStatus.FAILED,
+            ExecutionStatus.COMPLETED,
+        ]:
+            raise ValueError(
+                f"The execution of pipeline run {pipeline_run_id} has already "
+                "concluded and API tokens can no longer be generated for it "
+                "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."
+            )
+
+        if step_run.status in [
+            ExecutionStatus.FAILED,
+            ExecutionStatus.COMPLETED,
+        ]:
+            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."
+            )
+
+    return generate_access_token(
+        user_id=token.user_id,
+        # Keep the original API key and device token scopes
+        api_key=auth_context.api_key,
+        device=auth_context.device,
+        schedule_id=schedule_id,
+        pipeline_run_id=pipeline_run_id,
+        step_run_id=step_run_id,
+    ).access_token

zenml/zen_server/routers/logs_endpoints.py

@@ -0,0 +1,66 @@
+#  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.
+"""Endpoint definitions for logs."""
+
+from uuid import UUID
+
+from fastapi import APIRouter, Security
+
+from zenml.constants import (
+    API,
+    LOGS,
+    VERSION_1,
+)
+from zenml.models.v2.core.logs import LogsResponse
+from zenml.zen_server.auth import AuthContext, authorize
+from zenml.zen_server.exceptions import error_response
+from zenml.zen_server.rbac.endpoint_utils import (
+    verify_permissions_and_get_entity,
+)
+from zenml.zen_server.utils import (
+    handle_exceptions,
+    zen_store,
+)
+
+router = APIRouter(
+    prefix=API + VERSION_1 + LOGS,
+    tags=["logs"],
+    responses={401: error_response, 403: error_response},
+)
+
+
+@router.get(
+    "/{logs_id}",
+    response_model=LogsResponse,
+    responses={401: error_response, 404: error_response, 422: error_response},
+)
+@handle_exceptions
+def get_logs(
+    logs_id: UUID,
+    hydrate: bool = True,
+    _: AuthContext = Security(authorize),
+) -> LogsResponse:
+    """Returns the requested logs.
+
+    Args:
+        logs_id: ID of the logs.
+        hydrate: Flag deciding whether to hydrate the output model(s)
+            by including metadata fields in the response.
+
+    Returns:
+        The requested logs.
+    """
+    return verify_permissions_and_get_entity(
+        id=logs_id, get_method=zen_store().get_logs, hydrate=hydrate
+    )

zenml/zen_server/routers/workspaces_endpoints.py

@@ -74,7 +74,6 @@ from zenml.models import (
     PipelineRunRequest,
     PipelineRunResponse,
     RunMetadataRequest,
-    RunMetadataResponse,
     RunTemplateFilter,
     RunTemplateRequest,
     RunTemplateResponse,
@@ -977,7 +976,6 @@ def get_or_create_pipeline_run(
 
 @router.post(
     WORKSPACES + "/{workspace_name_or_id}" + RUN_METADATA,
-    response_model=List[RunMetadataResponse],
     responses={401: error_response, 409: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -985,7 +983,7 @@ def create_run_metadata(
     workspace_name_or_id: Union[str, UUID],
     run_metadata: RunMetadataRequest,
     auth_context: AuthContext = Security(authorize),
-) -> List[RunMetadataResponse]:
+) -> None:
     """Creates run metadata.
 
     Args:
@@ -1039,7 +1037,8 @@ def create_run_metadata(
         resource_type=ResourceType.RUN_METADATA, action=Action.CREATE
     )
 
-    return zen_store().create_run_metadata(run_metadata)
+    zen_store().create_run_metadata(run_metadata)
+    return None
 
 
 @router.post(

zenml/zen_server/template_execution/utils.py

@@ -42,7 +42,7 @@ from zenml.pipelines.run_utils import (
 )
 from zenml.stack.flavor import Flavor
 from zenml.utils import dict_utils, requirements_utils, settings_utils
-from zenml.zen_server.auth import AuthContext
+from zenml.zen_server.auth import AuthContext, generate_access_token
 from zenml.zen_server.template_execution.runner_entrypoint_configuration import (
     RunnerEntrypointConfiguration,
 )
@@ -111,17 +111,6 @@ def run_template(
 
     new_deployment = zen_store().create_deployment(deployment_request)
 
-    if auth_context.access_token:
-        token = auth_context.access_token
-        token.pipeline_id = deployment_request.pipeline
-
-        # We create a non-expiring token to make sure its active for the entire
-        # duration of the pipeline run
-        api_token = token.encode(expires=None)
-    else:
-        assert auth_context.encoded_access_token
-        api_token = auth_context.encoded_access_token
-
     server_url = server_config().server_url
     if not server_url:
         raise RuntimeError(
@@ -130,6 +119,18 @@ def run_template(
     assert build.zenml_version
     zenml_version = build.zenml_version
 
+    placeholder_run = create_placeholder_run(deployment=new_deployment)
+    assert placeholder_run
+
+    # We create an API token scoped to the pipeline run
+    api_token = generate_access_token(
+        user_id=auth_context.user.id,
+        pipeline_run_id=placeholder_run.id,
+        # Keep the original API key or device scopes, if any
+        api_key=auth_context.api_key,
+        device=auth_context.device,
+    ).access_token
+
     environment = {
         ENV_ZENML_ACTIVE_WORKSPACE_ID: str(new_deployment.workspace.id),
         ENV_ZENML_ACTIVE_STACK_ID: str(stack.id),
@@ -145,9 +146,6 @@ def run_template(
         deployment_id=new_deployment.id
     )
 
-    placeholder_run = create_placeholder_run(deployment=new_deployment)
-    assert placeholder_run
-
     def _task() -> None:
         pypi_requirements, apt_packages = (
             requirements_utils.get_requirements_for_stack(stack=stack)
@@ -350,7 +348,7 @@ def deployment_request_from_template(
     )
 
     step_config_dict_base = pipeline_configuration.model_dump(
-        exclude={"name", "parameters"}
+        exclude={"name", "parameters", "tags"}
     )
     steps = {}
     for invocation_id, step in deployment.step_configurations.items():

zenml/zen_server/utils.py

@@ -43,6 +43,7 @@ from zenml.enums import StoreType
 from zenml.exceptions import IllegalOperationError, OAuthError
 from zenml.logger import get_logger
 from zenml.plugins.plugin_flavor_registry import PluginFlavorRegistry
+from zenml.zen_server.cache import MemoryCache
 from zenml.zen_server.deploy.deployment import (
     LocalServerDeployment,
 )
@@ -67,6 +68,7 @@ _rbac: Optional[RBACInterface] = None
 _feature_gate: Optional[FeatureGateInterface] = None
 _workload_manager: Optional[WorkloadManagerInterface] = None
 _plugin_flavor_registry: Optional[PluginFlavorRegistry] = None
+_memcache: Optional[MemoryCache] = None
 
 
 def zen_store() -> "SqlZenStore":
@@ -222,6 +224,31 @@ def initialize_zen_store() -> None:
     _zen_store = zen_store_
 
 
+def initialize_memcache(max_capacity: int, default_expiry: int) -> None:
+    """Initialize the memory cache.
+
+    Args:
+        max_capacity: The maximum capacity of the cache.
+        default_expiry: The default expiry time in seconds.
+    """
+    global _memcache
+    _memcache = MemoryCache(max_capacity, default_expiry)
+
+
+def memcache() -> MemoryCache:
+    """Return the memory cache.
+
+    Returns:
+        The memory cache.
+
+    Raises:
+        RuntimeError: If the memory cache is not initialized.
+    """
+    if _memcache is None:
+        raise RuntimeError("Memory cache not initialized")
+    return _memcache
+
+
 _server_config: Optional[ServerConfiguration] = None
 
 

zenml/zen_server/zen_server_api.py

@@ -64,13 +64,13 @@ from zenml.zen_server.routers import (
     devices_endpoints,
     event_source_endpoints,
     flavors_endpoints,
+    logs_endpoints,
     model_versions_endpoints,
     models_endpoints,
     pipeline_builds_endpoints,
     pipeline_deployments_endpoints,
     pipelines_endpoints,
     plugin_endpoints,
-    run_metadata_endpoints,
     run_templates_endpoints,
     runs_endpoints,
     schedule_endpoints,
@@ -95,6 +95,7 @@ from zenml.zen_server.secure_headers import (
 )
 from zenml.zen_server.utils import (
     initialize_feature_gate,
+    initialize_memcache,
     initialize_plugins,
     initialize_rbac,
     initialize_workload_manager,
@@ -335,9 +336,10 @@ async def infer_source_context(request: Request, call_next: Any) -> Any:
 @app.on_event("startup")
 def initialize() -> None:
     """Initialize the ZenML server."""
+    cfg = server_config()
     # Set the maximum number of worker threads
     to_thread.current_default_thread_limiter().total_tokens = (
-        server_config().thread_pool_size
+        cfg.thread_pool_size
     )
     # IMPORTANT: these need to be run before the fastapi app starts, to avoid
     # race conditions
@@ -347,6 +349,7 @@ def initialize() -> None:
     initialize_workload_manager()
     initialize_plugins()
     initialize_secure_headers()
+    initialize_memcache(cfg.memcache_max_capacity, cfg.memcache_default_expiry)
 
 
 DASHBOARD_REDIRECT_URL = None
@@ -415,6 +418,7 @@ app.include_router(code_repositories_endpoints.router)
 app.include_router(plugin_endpoints.plugin_router)
 app.include_router(event_source_endpoints.event_source_router)
 app.include_router(flavors_endpoints.router)
+app.include_router(logs_endpoints.router)
 app.include_router(models_endpoints.router)
 app.include_router(model_versions_endpoints.router)
 app.include_router(model_versions_endpoints.model_version_artifacts_router)
@@ -423,7 +427,6 @@ app.include_router(pipelines_endpoints.router)
 app.include_router(pipeline_builds_endpoints.router)
 app.include_router(pipeline_deployments_endpoints.router)
 app.include_router(runs_endpoints.router)
-app.include_router(run_metadata_endpoints.router)
 app.include_router(run_templates_endpoints.router)
 app.include_router(schedule_endpoints.router)
 app.include_router(secrets_endpoints.router)

zenml/zen_stores/migrations/versions/0.70.0_release.py

@@ -0,0 +1,23 @@
+"""Release [0.70.0].
+
+Revision ID: 0.70.0
+Revises: 904464ea4041
+Create Date: 2024-11-12 12:20:05.537646
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "0.70.0"
+down_revision = "904464ea4041"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    pass
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    pass