zenml-nightly 0.58.1.dev20240607__tar.gz → 0.58.1.dev20240609__tar.gz

{zenml_nightly-0.58.1.dev20240607 → zenml_nightly-0.58.1.dev20240609}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zenml-nightly
-Version: 0.58.1.dev20240607
+Version: 0.58.1.dev20240609
 Summary: ZenML: Write production-ready ML code.
 Home-page: https://zenml.io
 License: Apache-2.0

{zenml_nightly-0.58.1.dev20240607 → zenml_nightly-0.58.1.dev20240609}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zenml-nightly"
-version = "0.58.1.dev20240607"
+version = "0.58.1.dev20240609"
 packages = [{ include = "zenml", from = "src" }]
 description = "ZenML: Write production-ready ML code."
 authors = ["ZenML GmbH <info@zenml.io>"]

zenml_nightly-0.58.1.dev20240609/src/zenml/VERSION

@@ -0,0 +1 @@
+0.58.1.dev20240609

{zenml_nightly-0.58.1.dev20240607 → zenml_nightly-0.58.1.dev20240609}/src/zenml/artifacts/utils.py

@@ -745,6 +745,8 @@ def _load_file_from_artifact_store(
     uri: str,
     artifact_store: "BaseArtifactStore",
     mode: str = "rb",
+    offset: int = 0,
+    length: Optional[int] = None,
 ) -> Any:
     """Load the given uri from the given artifact store.
 
@@ -752,6 +754,8 @@ def _load_file_from_artifact_store(
         uri: The uri of the file to load.
         artifact_store: The artifact store from which to load the file.
         mode: The mode in which to open the file.
+        offset: The offset from which to start reading.
+        length: The amount of bytes that should be read.
 
     Returns:
         The loaded file.
@@ -763,7 +767,19 @@ def _load_file_from_artifact_store(
     """
     try:
         with artifact_store.open(uri, mode) as text_file:
-            return text_file.read()
+            if offset < 0:
+                # If the offset is negative, we seek backwards from the end of
+                # the file
+                try:
+                    text_file.seek(offset, os.SEEK_END)
+                except OSError:
+                    # The negative offset was too large for the file, we seek
+                    # to the start of the file
+                    text_file.seek(0, os.SEEK_SET)
+            elif offset > 0:
+                text_file.seek(offset, os.SEEK_SET)
+
+            return text_file.read(length)
     except FileNotFoundError:
         raise DoesNotExistException(
             f"File '{uri}' does not exist in artifact store "

{zenml_nightly-0.58.1.dev20240607 → zenml_nightly-0.58.1.dev20240609}/src/zenml/zen_server/auth.py

@@ -238,8 +238,8 @@ def authenticate_credentials(
             decoded_token = JWTToken.decode_token(
                 token=access_token,
             )
-        except AuthorizationException:
-            error = "Authentication error: error decoding access token. You may need to rerun zenml connect."
+        except AuthorizationException as e:
+            error = f"Authentication error: error decoding access token: {e}."
             logger.exception(error)
             raise AuthorizationException(error)
 

{zenml_nightly-0.58.1.dev20240607 → zenml_nightly-0.58.1.dev20240609}/src/zenml/zen_server/jwt.py

@@ -89,7 +89,7 @@ class JWTToken(BaseModel):
         except jwt.PyJWTError as e:
             raise AuthorizationException(f"Invalid JWT token: {e}") from e
 
-        subject: str = claims.get("sub", "")
+        subject: str = claims.pop("sub", "")
         if not subject:
             raise AuthorizationException(
                 "Invalid JWT token: the subject claim is missing"
@@ -105,7 +105,7 @@ class JWTToken(BaseModel):
         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(
                     "Invalid JWT token: the device_id claim is not a valid "
@@ -115,7 +115,7 @@ 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(
                     "Invalid JWT token: the api_key_id claim is not a valid "
@@ -125,7 +125,7 @@ class JWTToken(BaseModel):
         pipeline_id: Optional[UUID] = None
         if "pipeline_id" in claims:
             try:
-                pipeline_id = UUID(claims["pipeline_id"])
+                pipeline_id = UUID(claims.pop("pipeline_id"))
             except ValueError:
                 raise AuthorizationException(
                     "Invalid JWT token: the pipeline_id claim is not a valid "
@@ -135,7 +135,7 @@ class JWTToken(BaseModel):
         schedule_id: Optional[UUID] = None
         if "schedule_id" in claims:
             try:
-                schedule_id = UUID(claims["schedule_id"])
+                schedule_id = UUID(claims.pop("schedule_id"))
             except ValueError:
                 raise AuthorizationException(
                     "Invalid JWT token: the schedule_id claim is not a valid "
@@ -165,14 +165,17 @@ 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:
@@ -182,9 +185,6 @@ class JWTToken(BaseModel):
         if self.schedule_id:
             claims["schedule_id"] = str(self.schedule_id)
 
-        # Apply custom claims
-        claims.update(self.claims)
-
         return jwt.encode(
             claims,
             config.jwt_secret_key,

{zenml_nightly-0.58.1.dev20240607 → zenml_nightly-0.58.1.dev20240609}/src/zenml/zen_server/routers/auth_endpoints.py

@@ -44,6 +44,7 @@ from zenml.enums import (
     OAuthDeviceStatus,
     OAuthGrantTypes,
 )
+from zenml.exceptions import AuthorizationException
 from zenml.logger import get_logger
 from zenml.models import (
     APIKeyInternalResponse,
@@ -510,6 +511,8 @@ def 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.
     """
     token = auth_context.access_token
     if not token or not auth_context.encoded_access_token:
@@ -523,6 +526,20 @@ def api_token(
         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}."
+        )
+
+    if schedule_id and token.schedule_id and schedule_id != token.schedule_id:
+        raise AuthorizationException(
+            f"Unable to scope API token to schedule {schedule_id}. The "
+            f"token used to authorize this request is already scoped to "
+            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

{zenml_nightly-0.58.1.dev20240607 → zenml_nightly-0.58.1.dev20240609}/src/zenml/zen_server/routers/steps_endpoints.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for steps (and artifacts) of pipeline runs."""
 
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 from uuid import UUID
 
 from fastapi import APIRouter, Depends, HTTPException, Security
@@ -239,12 +239,16 @@ def get_step_status(
 @handle_exceptions
 def get_step_logs(
     step_id: UUID,
+    offset: int = 0,
+    length: Optional[int] = 1024 * 1024 * 16,  # Default to 16MiB of data
     _: AuthContext = Security(authorize),
 ) -> str:
     """Get the logs of a specific step.
 
     Args:
         step_id: ID of the step for which to get the logs.
+        offset: The offset from which to start reading.
+        length: The amount of bytes that should be read.
 
     Returns:
         The logs of the step.
@@ -265,6 +269,10 @@ def get_step_logs(
     artifact_store = _load_artifact_store(logs.artifact_store_id, store)
     return str(
         _load_file_from_artifact_store(
-            logs.uri, artifact_store=artifact_store, mode="r"
-        )
+            logs.uri,
+            artifact_store=artifact_store,
+            mode="rb",
+            offset=offset,
+            length=length,
+        ).decode()
     )

{zenml_nightly-0.58.1.dev20240607 → zenml_nightly-0.58.1.dev20240609}/src/zenml/zen_stores/rest_zen_store.py

@@ -3730,7 +3730,13 @@ class RestZenStore(BaseZenStore):
         return self._session
 
     def clear_session(self) -> None:
-        """Clear the authentication session and any cached API tokens."""
+        """Clear the authentication session and any cached API tokens.
+
+        Raises:
+            AuthorizationException: If the API token can't be reset because
+                the store configuration does not contain username and password
+                or an API key to fetch a new token.
+        """
         self._session = None
         self._api_token = None
         # Clear the configured API token only if it's possible to fetch a new
@@ -3742,6 +3748,16 @@ class RestZenStore(BaseZenStore):
             or self.config.api_key is not None
         ):
             self.config.api_token = None
+        elif self.config.api_token:
+            raise AuthorizationException(
+                "Unable to refresh invalid API token. This is probably "
+                "because you're connected to your ZenML server with device "
+                "authentication. Rerunning `zenml connect --url "
+                f"{self.config.url}` should solve this issue. "
+                "If you're seeing this error from an automated workload, "
+                "you should probably use a service account to start that "
+                "workload to prevent this error."
+            )
 
     @staticmethod
     def _handle_response(response: requests.Response) -> Json:

zenml_nightly-0.58.1.dev20240607/src/zenml/VERSION

@@ -1 +0,0 @@
-0.58.1.dev20240607