zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/zen_server/feature_gate/zenml_cloud_feature_gate.py

@@ -0,0 +1,125 @@
+#  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.
+"""ZenML Pro implementation of the feature gate."""
+
+from typing import Any, Dict
+from uuid import UUID
+
+from pydantic import BaseModel, Field
+
+from zenml.config.server_config import ServerConfiguration
+from zenml.exceptions import SubscriptionUpgradeRequiredError
+from zenml.logger import get_logger
+from zenml.zen_server.cloud_utils import cloud_connection
+from zenml.zen_server.feature_gate.feature_gate_interface import (
+    FeatureGateInterface,
+)
+from zenml.zen_server.rbac.models import ResourceType
+
+logger = get_logger(__name__)
+
+server_config = ServerConfiguration.get_server_config()
+
+ORGANIZATION_ID = server_config.metadata.get("organization_id", "unknown")
+
+USAGE_EVENT_ENDPOINT = "/usage-event"
+ENTITLEMENT_ENDPOINT = f"/organizations/{ORGANIZATION_ID}/entitlement"
+
+
+class RawUsageEvent(BaseModel):
+    """Model for reporting raw usage of a feature.
+
+    In case of consumables the UsageReport allows the Pricing Backend to
+    increment the usage per time-frame by 1.
+    """
+
+    organization_id: str = Field(
+        description="The organization that this usage can be attributed to.",
+    )
+    feature: ResourceType = Field(
+        description="The feature whose usage is being reported.",
+    )
+    total: int = Field(
+        description="The total amount of entities of this type."
+    )
+    metadata: Dict[str, Any] = Field(
+        default={},
+        description="Allows attaching additional metadata to events.",
+    )
+
+
+class ZenMLCloudFeatureGateInterface(FeatureGateInterface):
+    """ZenML Cloud Feature Gate implementation."""
+
+    def __init__(self) -> None:
+        """Initialize the object."""
+        self._connection = cloud_connection()
+
+    def check_entitlement(self, resource: ResourceType) -> None:
+        """Checks if a user is entitled to create a resource.
+
+        Args:
+            resource: The resource the user wants to create
+
+        Raises:
+            SubscriptionUpgradeRequiredError: in case a subscription limit is reached
+        """
+        try:
+            response = self._connection.get(
+                endpoint=ENTITLEMENT_ENDPOINT + "/" + resource, params=None
+            )
+        except SubscriptionUpgradeRequiredError:
+            raise SubscriptionUpgradeRequiredError(
+                f"Your subscription reached its `{resource}` limit. Please "
+                f"upgrade your subscription or reach out to us."
+            )
+
+        if response.status_code != 200:
+            logger.warning(
+                "Unexpected response status code from entitlement "
+                f"endpoint: {response.status_code}. Message: "
+                f"{response.json()}"
+            )
+
+    def report_event(
+        self,
+        resource: ResourceType,
+        resource_id: UUID,
+        is_decrement: bool = False,
+    ) -> None:
+        """Reports the usage of a feature to the aggregator backend.
+
+        Args:
+            resource: The resource the user created
+            resource_id: ID of the resource that was created/deleted.
+            is_decrement: In case this event reports an actual decrement of usage
+        """
+        data = RawUsageEvent(
+            organization_id=ORGANIZATION_ID,
+            feature=resource,
+            total=1 if not is_decrement else -1,
+            metadata={
+                "tenant_id": str(server_config.get_external_server_id()),
+                "resource_id": str(resource_id),
+            },
+        ).model_dump()
+        response = self._connection.post(
+            endpoint=USAGE_EVENT_ENDPOINT, data=data
+        )
+        if response.status_code != 200:
+            logger.error(
+                "Usage report not accepted by upstream backend. "
+                f"Status Code: {response.status_code}, Message: "
+                f"{response.json()}."
+            )

zenml/zen_server/jwt.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.
-"""Authentication module for ZenML server."""
+"""JWT utilities module for ZenML server."""
 
 from datetime import datetime, timedelta
 from typing import (
@@ -25,7 +25,7 @@ from uuid import UUID
 import jwt
 from pydantic import BaseModel
 
-from zenml.exceptions import AuthorizationException
+from zenml.exceptions import CredentialsNotValid
 from zenml.logger import get_logger
 from zenml.zen_server.utils import server_config
 
@@ -40,16 +40,22 @@ class JWTToken(BaseModel):
         device_id: The id of the authenticated device.
         api_key_id: The id of the authenticated API key for which this token
             was issued.
-        pipeline_id: The id of the pipeline for which the token was issued.
         schedule_id: The id of the schedule for which the token was issued.
+        pipeline_run_id: The id of the pipeline run for which the token was
+            issued.
+        step_run_id: The id of the step run for which the token was
+            issued.
+        session_id: The id of the authenticated session (used for CSRF).
         claims: The original token claims.
     """
 
     user_id: UUID
     device_id: Optional[UUID] = None
     api_key_id: Optional[UUID] = None
-    pipeline_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] = {}
 
     @classmethod
@@ -71,7 +77,7 @@ class JWTToken(BaseModel):
             The decoded JWT access token.
 
         Raises:
-            AuthorizationException: If the token is invalid.
+            CredentialsNotValid: If the token is invalid.
         """
         config = server_config()
 
@@ -87,27 +93,27 @@ class JWTToken(BaseModel):
             )
             claims = cast(Dict[str, Any], claims_data)
         except jwt.PyJWTError as e:
-            raise AuthorizationException(f"Invalid JWT token: {e}") from e
+            raise CredentialsNotValid(f"Invalid JWT token: {e}") from e
 
-        subject: str = claims.get("sub", "")
+        subject: str = claims.pop("sub", "")
         if not subject:
-            raise AuthorizationException(
+            raise CredentialsNotValid(
                 "Invalid JWT token: the subject claim is missing"
             )
 
         try:
             user_id = UUID(subject)
         except ValueError:
-            raise AuthorizationException(
+            raise CredentialsNotValid(
                 "Invalid JWT token: the subject claim is not a valid UUID"
             )
 
         device_id: Optional[UUID] = None
         if "device_id" in claims:
             try:
-                device_id = UUID(claims["device_id"])
+                device_id = UUID(claims.pop("device_id"))
             except ValueError:
-                raise AuthorizationException(
+                raise CredentialsNotValid(
                     "Invalid JWT token: the device_id claim is not a valid "
                     "UUID"
                 )
@@ -115,30 +121,50 @@ class JWTToken(BaseModel):
         api_key_id: Optional[UUID] = None
         if "api_key_id" in claims:
             try:
-                api_key_id = UUID(claims["api_key_id"])
+                api_key_id = UUID(claims.pop("api_key_id"))
             except ValueError:
-                raise AuthorizationException(
+                raise CredentialsNotValid(
                     "Invalid JWT token: the api_key_id claim is not a valid "
                     "UUID"
                 )
 
-        pipeline_id: Optional[UUID] = None
-        if "pipeline_id" in claims:
+        schedule_id: Optional[UUID] = None
+        if "schedule_id" in claims:
             try:
-                pipeline_id = UUID(claims["pipeline_id"])
+                schedule_id = UUID(claims.pop("schedule_id"))
             except ValueError:
-                raise AuthorizationException(
-                    "Invalid JWT token: the pipeline_id claim is not a valid "
+                raise CredentialsNotValid(
+                    "Invalid JWT token: the schedule_id claim is not a valid "
                     "UUID"
                 )
 
-        schedule_id: Optional[UUID] = None
-        if "schedule_id" in claims:
+        pipeline_run_id: Optional[UUID] = None
+        if "pipeline_run_id" in claims:
             try:
-                schedule_id = UUID(claims["schedule_id"])
+                pipeline_run_id = UUID(claims.pop("pipeline_run_id"))
             except ValueError:
-                raise AuthorizationException(
-                    "Invalid JWT token: the schedule_id claim is not a valid "
+                raise CredentialsNotValid(
+                    "Invalid JWT token: the pipeline_run_id claim is not a valid "
+                    "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:
+                session_id = UUID(claims.pop("session_id"))
+            except ValueError:
+                raise CredentialsNotValid(
+                    "Invalid JWT token: the session_id claim is not a valid "
                     "UUID"
                 )
 
@@ -146,8 +172,10 @@ class JWTToken(BaseModel):
             user_id=user_id,
             device_id=device_id,
             api_key_id=api_key_id,
-            pipeline_id=pipeline_id,
             schedule_id=schedule_id,
+            pipeline_run_id=pipeline_run_id,
+            step_run_id=step_run_id,
+            session_id=session_id,
             claims=claims,
         )
 
@@ -165,25 +193,29 @@ class JWTToken(BaseModel):
         """
         config = server_config()
 
-        claims: Dict[str, Any] = dict(
-            sub=str(self.user_id),
-        )
+        claims: Dict[str, Any] = self.claims.copy()
+
+        claims["sub"] = str(self.user_id)
         claims["iss"] = config.get_jwt_token_issuer()
         claims["aud"] = config.get_jwt_token_audience()
 
         if expires:
             claims["exp"] = expires
+        else:
+            claims.pop("exp", None)
+
         if self.device_id:
             claims["device_id"] = str(self.device_id)
         if self.api_key_id:
             claims["api_key_id"] = str(self.api_key_id)
-        if self.pipeline_id:
-            claims["pipeline_id"] = str(self.pipeline_id)
         if self.schedule_id:
             claims["schedule_id"] = str(self.schedule_id)
-
-        # Apply custom claims
-        claims.update(self.claims)
+        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)
 
         return jwt.encode(
             claims,

zenml/zen_server/rate_limit.py

@@ -0,0 +1,200 @@
+#  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.
+"""Rate limiting for the ZenML Server."""
+
+import inspect
+import time
+from collections import defaultdict
+from contextlib import contextmanager
+from functools import wraps
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Generator,
+    List,
+    Optional,
+    TypeVar,
+    cast,
+)
+
+from starlette.requests import Request
+
+from zenml.logger import get_logger
+from zenml.zen_server.utils import server_config
+
+logger = get_logger(__name__)
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+class RequestLimiter:
+    """Simple in-memory rate limiter."""
+
+    def __init__(
+        self,
+        day_limit: Optional[int] = None,
+        minute_limit: Optional[int] = None,
+    ):
+        """Initializes the limiter.
+
+        Args:
+            day_limit: The number of requests allowed per day.
+            minute_limit: The number of requests allowed per minute.
+
+        Raises:
+            ValueError: If both day_limit and minute_limit are None.
+        """
+        self.limiting_enabled = server_config().rate_limit_enabled
+        if not self.limiting_enabled:
+            return
+        if day_limit is None and minute_limit is None:
+            raise ValueError("Pass either day or minuter limits, or both.")
+        self.day_limit = day_limit
+        self.minute_limit = minute_limit
+        self.limiter: Dict[str, List[float]] = defaultdict(list)
+
+    def hit_limiter(self, request: Request) -> None:
+        """Increase the number of hits in the limiter.
+
+        Args:
+            request: Request object.
+
+        Raises:
+            HTTPException: If the request limit is exceeded.
+        """
+        if not self.limiting_enabled:
+            return
+        from fastapi import HTTPException
+
+        requester = self._get_ipaddr(request)
+        now = time.time()
+        minute_ago = now - 60
+        day_ago = now - 60 * 60 * 24
+        self.limiter[requester].append(now)
+
+        from bisect import bisect_left
+
+        # remove failures older than a day
+        older_index = bisect_left(self.limiter[requester], day_ago)
+        self.limiter[requester] = self.limiter[requester][older_index:]
+
+        if self.day_limit and len(self.limiter[requester]) > self.day_limit:
+            raise HTTPException(
+                status_code=429, detail="Daily request limit exceeded."
+            )
+        minute_requests = len(
+            [
+                limiter_hit
+                for limiter_hit in self.limiter[requester][::-1]
+                if limiter_hit >= minute_ago
+            ]
+        )
+        if self.minute_limit and minute_requests > self.minute_limit:
+            raise HTTPException(
+                status_code=429, detail="Minute request limit exceeded."
+            )
+
+    def reset_limiter(self, request: Request) -> None:
+        """Resets the limiter on successful request.
+
+        Args:
+            request: Request object.
+        """
+        if self.limiting_enabled:
+            requester = self._get_ipaddr(request)
+            if requester in self.limiter:
+                del self.limiter[requester]
+
+    def _get_ipaddr(self, request: Request) -> str:
+        """Returns the IP address for the current request.
+
+        Based on the X-Forwarded-For headers or client information.
+
+        Args:
+            request: The request object.
+
+        Returns:
+            The ip address for the current request (or 127.0.0.1 if none found).
+        """
+        if "X_FORWARDED_FOR" in request.headers:
+            return request.headers["X_FORWARDED_FOR"]
+        else:
+            if not request.client or not request.client.host:
+                return "127.0.0.1"
+
+            return request.client.host
+
+    @contextmanager
+    def limit_failed_requests(
+        self, request: Request
+    ) -> Generator[None, Any, Any]:
+        """Limits the number of failed requests.
+
+        Args:
+            request: Request object.
+
+        Yields:
+            None
+        """
+        self.hit_limiter(request)
+
+        yield
+
+        # if request was successful - reset limiter
+        self.reset_limiter(request)
+
+
+def rate_limit_requests(
+    day_limit: Optional[int] = None,
+    minute_limit: Optional[int] = None,
+) -> Callable[..., Any]:
+    """Decorator to handle exceptions in the API.
+
+    Args:
+        day_limit: Number of requests allowed per day.
+        minute_limit: Number of requests allowed per minute.
+
+    Returns:
+        Decorated function.
+    """
+    limiter = RequestLimiter(day_limit=day_limit, minute_limit=minute_limit)
+
+    def decorator(func: F) -> F:
+        request_arg, request_kwarg = None, None
+        parameters = inspect.signature(func).parameters
+        for arg_num, arg_name in enumerate(parameters):
+            if parameters[arg_name].annotation == Request:
+                request_arg = arg_num
+                request_kwarg = arg_name
+                break
+        if request_arg is None or request_kwarg is None:
+            raise ValueError(
+                "Rate limiting APIs must have argument of `Request` type."
+            )
+
+        @wraps(func)
+        def decorated(
+            *args: Any,
+            **kwargs: Any,
+        ) -> Any:
+            if request_kwarg in kwargs:
+                request = kwargs[request_kwarg]
+            else:
+                request = args[request_arg]
+            with limiter.limit_failed_requests(request):
+                return func(*args, **kwargs)
+
+        return cast(F, decorated)
+
+    return decorator

zenml/zen_server/rbac/endpoint_utils.py

@@ -1,18 +1,39 @@
+#  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.
 """High-level helper functions to write endpoints with RBAC."""
-from typing import Any, Callable, TypeVar, Union
+
+from typing import Any, Callable, List, TypeVar, Union
 from uuid import UUID
 
 from pydantic import BaseModel
 
+from zenml.constants import (
+    REQUIRES_CUSTOM_RESOURCE_REPORTING,
+)
 from zenml.exceptions import IllegalOperationError
 from zenml.models import (
     BaseFilter,
+    BaseIdentifiedResponse,
     BaseRequest,
-    BaseResponse,
     Page,
     UserScopedRequest,
 )
 from zenml.zen_server.auth import get_auth_context
+from zenml.zen_server.feature_gate.endpoint_utils import (
+    check_entitlement,
+    report_usage,
+)
 from zenml.zen_server.rbac.models import Action, ResourceType
 from zenml.zen_server.rbac.utils import (
     dehydrate_page,
@@ -21,9 +42,10 @@ from zenml.zen_server.rbac.utils import (
     verify_permission,
     verify_permission_for_model,
 )
+from zenml.zen_server.utils import server_config
 
 AnyRequest = TypeVar("AnyRequest", bound=BaseRequest)
-AnyResponse = TypeVar("AnyResponse", bound=BaseResponse)  # type: ignore[type-arg]
+AnyResponse = TypeVar("AnyResponse", bound=BaseIdentifiedResponse)  # type: ignore[type-arg]
 AnyFilter = TypeVar("AnyFilter", bound=BaseFilter)
 AnyUpdate = TypeVar("AnyUpdate", bound=BaseModel)
 UUIDOrStr = TypeVar("UUIDOrStr", UUID, Union[UUID, str])
@@ -57,9 +79,63 @@ def verify_permissions_and_create_entity(
                 f"Not allowed to create resource '{resource_type}' for a "
                 "different user."
             )
+    verify_permission(resource_type=resource_type, action=Action.CREATE)
+
+    needs_usage_increment = (
+        resource_type in server_config().reportable_resources
+        and resource_type not in REQUIRES_CUSTOM_RESOURCE_REPORTING
+    )
+    if needs_usage_increment:
+        check_entitlement(resource_type)
+
+    created = create_method(request_model)
+
+    if needs_usage_increment:
+        report_usage(resource_type, resource_id=created.id)
+
+    return created
+
+
+def verify_permissions_and_batch_create_entity(
+    batch: List[AnyRequest],
+    resource_type: ResourceType,
+    create_method: Callable[[List[AnyRequest]], List[AnyResponse]],
+) -> List[AnyResponse]:
+    """Verify permissions and create a batch of entities if authorized.
+
+    Args:
+        batch: The batch to create.
+        resource_type: The resource type of the entities to create.
+        create_method: The method to create the entities.
+
+    Raises:
+        IllegalOperationError: If the request model has a different owner then
+            the currently authenticated user.
+        RuntimeError: If the resource type is usage-tracked.
+
+    Returns:
+        The created entities.
+    """
+    auth_context = get_auth_context()
+    assert auth_context
+
+    for request_model in batch:
+        if isinstance(request_model, UserScopedRequest):
+            if request_model.user != auth_context.user.id:
+                raise IllegalOperationError(
+                    f"Not allowed to create resource '{resource_type}' for a "
+                    "different user."
+                )
 
     verify_permission(resource_type=resource_type, action=Action.CREATE)
-    return create_method(request_model)
+
+    if resource_type in server_config().reportable_resources:
+        raise RuntimeError(
+            "Batch requests are currently not possible with usage-tracked features."
+        )
+
+    created = create_method(batch)
+    return created
 
 
 def verify_permissions_and_get_entity(
@@ -113,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.
@@ -127,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)
@@ -135,20 +212,26 @@ 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],
-) -> None:
+) -> AnyResponse:
     """Verify permissions and delete an entity.
 
     Args:
         id: The ID of the entity to delete.
         get_method: The method to fetch the entity.
         delete_method: The method to delete the 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)
 
+    return model
+
 
 def verify_permissions_and_prune_entities(
     resource_type: ResourceType,