zenml-nightly 0.58.1.dev20240607__py3-none-any.whl → 0.58.1.dev20240609__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.58.1.dev20240607
+0.58.1.dev20240609

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/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/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/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/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/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.dist-info → zenml_nightly-0.58.1.dev20240609.dist-info}/METADATA

@@ -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.dist-info → zenml_nightly-0.58.1.dev20240609.dist-info}/RECORD

@@ -6,7 +6,7 @@ RELEASE_NOTES.md,sha256=nXkCYUiBFMOQrSB9Vmsl5zouXPEBE90pdxDp75jxxwU,329317
 ROADMAP.md,sha256=hiLSmr16BH8Dfx7SaQM4JcXCGCVl6mFZPFAwJeDTrJU,407
 SECURITY.md,sha256=9DepA8y03yvCZLHEfcXLTDH4lUyKHquAdukBsccNN7c,682
 zenml/README.md,sha256=827dekbOWAs1BpW7VF1a4d7EbwPbjwccX-2zdXBENZo,1777
-zenml/VERSION,sha256=B76BYXykJsh0PGXOE7JJl6BC7xbpqxYAZTYkWSlIdHA,19
+zenml/VERSION,sha256=COl3SgNWbyFl1SMIf5AT0gOmTLSZ4VoLOlmyyr0G8Ao,19
 zenml/__init__.py,sha256=qT3QHiuTgUwk05sOiML0Gcjre3pugvDAw84_br9Psvo,2394
 zenml/_hub/__init__.py,sha256=6qDzpQAAZa__Aiiz0mC1qM-9dw9_jk_v_aXeJknxbDE,644
 zenml/_hub/client.py,sha256=LAxLuDjY9ya1PPRzxS0_pvu__FSpAqEjbtHTqVwgMWc,9071
@@ -36,7 +36,7 @@ zenml/artifacts/artifact_config.py,sha256=pG4qqnIRF5V9we_mIPIbwaynUpWAqOCunH8Ex5
 zenml/artifacts/external_artifact.py,sha256=IG5wKuPzy8nxDZyMb34C1TD81GJW0Ot_gtxQsSod8rY,6117
 zenml/artifacts/external_artifact_config.py,sha256=RSsu9ZBLgnI6YvB1Dxw9H_u66ZXcac28o2QuN6EABWI,3861
 zenml/artifacts/unmaterialized_artifact.py,sha256=tQX1ngze4vik4VTpSGN00ozyCv0ebHouULjTYtqOsf8,1514
-zenml/artifacts/utils.py,sha256=mabZRLgzyhFDJojqboSPWjt1ZwcijEaMpazsuwBpcW4,30470
+zenml/artifacts/utils.py,sha256=6H6e0a5e7LOEq4Po4GBj6GrXoCSETpN2DPZ7zlF14gE,31144
 zenml/cli/__init__.py,sha256=ZRjE6Ug536SseT7puiCyMmcvw4mr0P_ZBoa-RmwIiEQ,77081
 zenml/cli/annotator.py,sha256=tEdducGdFn57DFLJVZQ-MyXH1auTGFueRmDc78N-vPQ,6970
 zenml/cli/artifact.py,sha256=UY_inYCNCKPzHyFXfSRnpR2xHaQYzqL1hvKGxaRHKMU,9012
@@ -734,7 +734,7 @@ zenml/utils/uuid_utils.py,sha256=aOGQ2SdREexcVQICPU2jUAgjvAJxTmh4ESdM52PEhck,204
 zenml/utils/visualization_utils.py,sha256=SawUfw1OTkU2GV5rH7tw7hd3eqpFooiB4Fvr_7MoTus,4555
 zenml/utils/yaml_utils.py,sha256=DIZOD7iExpvcjDlInCp3hhKnx2o4vmVBsNZEQYQcN44,5750
 zenml/zen_server/__init__.py,sha256=WyltI9TzFW2mEHZVOs6alLWMCQrrZaFALtrQXs83STA,1355
-zenml/zen_server/auth.py,sha256=NxeRO43uXmFH0-s6B8aWsVeN6DUPvL-SqCKdAK2-0cg,26264
+zenml/zen_server/auth.py,sha256=KhPFDKpIMyQFl-yDacOljNZTrPc87S-RsqwWTq1urcM,26238
 zenml/zen_server/cloud_utils.py,sha256=w7pdc-9YJyqL7tvrGhTDINf2fum8HdtpRBZAUY9fiNg,6327
 zenml/zen_server/dashboard/assets/404-D5p6PIdn.js,sha256=8MB4WqhTpuaeEAwYhEjnzJKipxtK_0_Mor37gBsyIgE,1033
 zenml/zen_server/dashboard/assets/@radix-C9DBgJhe.js,sha256=MrmwKYEEzL1fZ5cWQ5TCAPTAjKZV5I-lVUEx_hbDFWc,262349
@@ -1053,7 +1053,7 @@ zenml/zen_server/feature_gate/__init__.py,sha256=yabe4fBY6NSusn-QlKQDLOvXVLERNpc
 zenml/zen_server/feature_gate/endpoint_utils.py,sha256=upQzEOlrD5c6cg0TfxYVM4BKDHle3LIV_c7NZHoaKvg,2220
 zenml/zen_server/feature_gate/feature_gate_interface.py,sha256=AdYAsHiNbGaloxSHpm7DLSsa8CDEGu8ikyxp9U3a0wU,1639
 zenml/zen_server/feature_gate/zenml_cloud_feature_gate.py,sha256=_qpA4qGgtCAwZX5stcj4auz5aWDVi-ouPiUqS20LFPo,4250
-zenml/zen_server/jwt.py,sha256=AcazqLIFGD_eSoR-toQcUJAdEIlpoSLbHI2zbGsWtes,6139
+zenml/zen_server/jwt.py,sha256=cMJS24EDXqhZ_-hOFDORsZMgNF6JcfSGNYFOuh99u_w,6151
 zenml/zen_server/pipeline_deployment/__init__.py,sha256=79knXLKfegsvVSVSWecpqrepq6iAavTUA4hKuiDk-WE,613
 zenml/zen_server/pipeline_deployment/runner_entrypoint_configuration.py,sha256=PpraJ77nCOUesVC6IBxNTeVMdS3eWG2vt-4GrT8hORQ,1682
 zenml/zen_server/pipeline_deployment/utils.py,sha256=T4dqRAl3Khb3UdbfPEpDf4Mek_mnE9FRbTJ_vmoj6Nw,13039
@@ -1068,7 +1068,7 @@ zenml/zen_server/rbac/zenml_cloud_rbac.py,sha256=_rWFCmZRk5Z5TzmWmCJAYzSoXSz1C0d
 zenml/zen_server/routers/__init__.py,sha256=ViyAhWL-ogHxE9wBvB_iMcur5H1NRVrzXkpogVY7FBA,641
 zenml/zen_server/routers/artifact_endpoint.py,sha256=XhbOat2pddDluiT0a_QH2RfNcqlwtf3yf2Fh92a6ZDw,5175
 zenml/zen_server/routers/artifact_version_endpoints.py,sha256=o9RVrnzuiY-AV6Nj10YwXQ_jy-OqCVMk1hEKfmihPEQ,7751
-zenml/zen_server/routers/auth_endpoints.py,sha256=U-Fsvi2oZAaYTj5LYMAokWDofam7YC9cZPzo4iAVWLY,18440
+zenml/zen_server/routers/auth_endpoints.py,sha256=xl8g5AVz_hiHVpofaLxyWHc2SDshH6028kYtHY6IRg0,19289
 zenml/zen_server/routers/code_repositories_endpoints.py,sha256=WFCRPsv3Qrm8QtCr5zxfSdgCS5WI1DPNCF4Y6vLXGow,4747
 zenml/zen_server/routers/devices_endpoints.py,sha256=2YoVb_a3fEjsYs8wMQYkzz9qM0n3ylHB6yTVNmhIVRU,10598
 zenml/zen_server/routers/event_source_endpoints.py,sha256=dXupWrySV3LtxsqMVoYpUJ-OK7q6o6ehfuYW_RU1JlA,10379
@@ -1089,7 +1089,7 @@ zenml/zen_server/routers/service_connectors_endpoints.py,sha256=Imf1WAKUiATPiQsi
 zenml/zen_server/routers/service_endpoints.py,sha256=PyAxQLewVnpotfQI_OmyUTl7ohTCrWrteA9nfvxDRyU,5055
 zenml/zen_server/routers/stack_components_endpoints.py,sha256=mFVeAZY2QmlEdeowvaFO0wy6mG2zQlybOBvrFg4zWdU,6110
 zenml/zen_server/routers/stacks_endpoints.py,sha256=imL7s26xFOf4EY4zwSz8y8-Ggl6-Xfu-tJY75TNodeo,4540
-zenml/zen_server/routers/steps_endpoints.py,sha256=kd3CFI2NXyaECSl9tlvdGCjt223tRiyvp2g96BjK_JA,7511
+zenml/zen_server/routers/steps_endpoints.py,sha256=oWI6B1erltAbG_PmVFjZeH2JElH7J46ecnp77lKO52s,7818
 zenml/zen_server/routers/tags_endpoints.py,sha256=oK-A-EqAsrIXXgl7A870I2PT8_fct1dZVQDQ-g8GHys,4941
 zenml/zen_server/routers/triggers_endpoints.py,sha256=FpnDKcWirLqV757radYCcEeNOUca3tjsmQnnLJMJJRs,14445
 zenml/zen_server/routers/users_endpoints.py,sha256=f-ZuuqtOwiyhtucJ3RAJ4IzeQVYcPTaQBca_wAtlxcU,25199
@@ -1245,7 +1245,7 @@ zenml/zen_stores/migrations/versions/ec0d785ca296_create_run_metadata_table.py,s
 zenml/zen_stores/migrations/versions/f3b3964e3a0f_add_oauth_devices.py,sha256=2CR4R-7Vx6j_AXxo-e5Guy6OX-ZnS47HSKSGfqlO-y0,3065
 zenml/zen_stores/migrations/versions/f49904a80aa7_increase_length_of_artifact_table_sources.py,sha256=kLgfDUnQdAb5_SyFx3VKXDLC0YbuBKf9iXRDNeBin7Q,1618
 zenml/zen_stores/migrations/versions/fbd7f18ced1e_increase_step_run_field_lengths.py,sha256=mqNSUNTkrzkGmOscmwmR98lqC7hFH4PXQPxTPFvBxps,1964
-zenml/zen_stores/rest_zen_store.py,sha256=zRTt6t9VWAu2B1NhP2RCENoIKsxIUt8ft8U0cuVdkxU,137354
+zenml/zen_stores/rest_zen_store.py,sha256=72twP2owz52SY6ecqiDLGX2t_vwAx2J1LLTBdtN0cs4,138159
 zenml/zen_stores/schemas/__init__.py,sha256=y9bkvHil9HkHU5s2eFRhnMlQM3n2Idu83kva-hjrOtM,4080
 zenml/zen_stores/schemas/api_key_schemas.py,sha256=lyFE3gtCQx00r4jaUaEUYcqcGWvjqpPdatb-WrOExPk,7064
 zenml/zen_stores/schemas/artifact_schemas.py,sha256=8YDgZwmWSnFrENyyzvNU4XYKW9Rg8_WdJHfU7AxGiYc,13170
@@ -1287,8 +1287,8 @@ zenml/zen_stores/secrets_stores/service_connector_secrets_store.py,sha256=0CgWJo
 zenml/zen_stores/secrets_stores/sql_secrets_store.py,sha256=Cips_SPX7CtfYQyOZNiHAhmVtVgb17HCyHQia5MlTX4,8653
 zenml/zen_stores/sql_zen_store.py,sha256=f5spF_vWJg8kO6ud9uzELqh7LRxSAWyB1HLono1_gT0,353540
 zenml/zen_stores/zen_store_interface.py,sha256=vi0tr6vDES97h_rFZZYeaz2Y0nM4xm66CcUXSOEITN4,86372
-zenml_nightly-0.58.1.dev20240607.dist-info/LICENSE,sha256=wbnfEnXnafPbqwANHkV6LUsPKOtdpsd-SNw37rogLtc,11359
-zenml_nightly-0.58.1.dev20240607.dist-info/METADATA,sha256=aM_bXFMT5x-uqu3-fUs-MowURaMqFu3o8CxG2Jxjm94,19094
-zenml_nightly-0.58.1.dev20240607.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-zenml_nightly-0.58.1.dev20240607.dist-info/entry_points.txt,sha256=QK3ETQE0YswAM2mWypNMOv8TLtr7EjnqAFq1br_jEFE,43
-zenml_nightly-0.58.1.dev20240607.dist-info/RECORD,,
+zenml_nightly-0.58.1.dev20240609.dist-info/LICENSE,sha256=wbnfEnXnafPbqwANHkV6LUsPKOtdpsd-SNw37rogLtc,11359
+zenml_nightly-0.58.1.dev20240609.dist-info/METADATA,sha256=cZ-_RALqyhbiBd4HYXdTm_iBJKPaBcLDkyzhPq2l9Vw,19094
+zenml_nightly-0.58.1.dev20240609.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+zenml_nightly-0.58.1.dev20240609.dist-info/entry_points.txt,sha256=QK3ETQE0YswAM2mWypNMOv8TLtr7EjnqAFq1br_jEFE,43
+zenml_nightly-0.58.1.dev20240609.dist-info/RECORD,,