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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (232) hide show
  1. README.md +1 -1
  2. RELEASE_NOTES.md +77 -0
  3. zenml/VERSION +1 -1
  4. zenml/__init__.py +2 -0
  5. zenml/artifacts/external_artifact.py +2 -1
  6. zenml/artifacts/utils.py +138 -79
  7. zenml/cli/base.py +4 -4
  8. zenml/cli/model.py +1 -6
  9. zenml/cli/stack.py +1 -0
  10. zenml/client.py +29 -74
  11. zenml/config/server_config.py +17 -1
  12. zenml/constants.py +2 -7
  13. zenml/data_validators/base_data_validator.py +2 -2
  14. zenml/enums.py +20 -4
  15. zenml/exceptions.py +4 -0
  16. zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py +20 -18
  17. zenml/integrations/azure/orchestrators/azureml_orchestrator.py +1 -1
  18. zenml/integrations/deepchecks/data_validators/deepchecks_data_validator.py +1 -1
  19. zenml/integrations/evidently/__init__.py +1 -1
  20. zenml/integrations/gcp/orchestrators/vertex_orchestrator.py +1 -1
  21. zenml/integrations/lightning/orchestrators/lightning_orchestrator.py +29 -9
  22. zenml/integrations/tensorboard/visualizers/tensorboard_visualizer.py +60 -54
  23. zenml/integrations/vllm/services/vllm_deployment.py +16 -7
  24. zenml/metadata/lazy_load.py +20 -7
  25. zenml/model/model.py +17 -64
  26. zenml/model/utils.py +5 -0
  27. zenml/models/__init__.py +0 -12
  28. zenml/models/v2/base/filter.py +121 -8
  29. zenml/models/v2/core/artifact_version.py +42 -7
  30. zenml/models/v2/core/model_version.py +26 -5
  31. zenml/models/v2/core/pipeline_run.py +25 -6
  32. zenml/models/v2/core/run_metadata.py +2 -217
  33. zenml/models/v2/core/step_run.py +62 -24
  34. zenml/orchestrators/base_orchestrator.py +12 -1
  35. zenml/orchestrators/input_utils.py +44 -19
  36. zenml/orchestrators/step_launcher.py +4 -3
  37. zenml/orchestrators/step_run_utils.py +19 -15
  38. zenml/orchestrators/step_runner.py +21 -13
  39. zenml/orchestrators/utils.py +45 -26
  40. zenml/stack_deployments/aws_stack_deployment.py +23 -6
  41. zenml/stack_deployments/azure_stack_deployment.py +28 -5
  42. zenml/stack_deployments/gcp_stack_deployment.py +25 -8
  43. zenml/stack_deployments/stack_deployment.py +3 -5
  44. zenml/steps/base_step.py +1 -1
  45. zenml/steps/entrypoint_function_utils.py +3 -5
  46. zenml/steps/step_context.py +3 -2
  47. zenml/steps/utils.py +13 -2
  48. zenml/utils/metadata_utils.py +335 -0
  49. zenml/zen_server/auth.py +221 -3
  50. zenml/zen_server/cache.py +208 -0
  51. zenml/zen_server/dashboard/assets/{404-DT4QRUqN.js → 404-NVXKFp-x.js} +1 -1
  52. zenml/zen_server/dashboard/assets/{@radix-DP6vWzyx.js → @radix-DeK6qiuw.js} +1 -1
  53. zenml/zen_server/dashboard/assets/{@react-router-BMhZulnd.js → @react-router-B3Z5rLr2.js} +1 -1
  54. zenml/zen_server/dashboard/assets/{@reactflow-8U9qNlMR.js → @reactflow-CK0KJUen.js} +2 -2
  55. zenml/zen_server/dashboard/assets/{@tanstack-BUCbhJyH.js → @tanstack-DT5WLu9C.js} +1 -1
  56. zenml/zen_server/dashboard/assets/AlertDialogDropdownItem-DezXKmDf.js +1 -0
  57. zenml/zen_server/dashboard/assets/{CodeSnippet-CqybNv0k.js → CodeSnippet-JzR8CEtw.js} +2 -2
  58. zenml/zen_server/dashboard/assets/{CollapsibleCard-0r_8G2Lj.js → CollapsibleCard-DQW_ktMO.js} +1 -1
  59. zenml/zen_server/dashboard/assets/{Commands-BDjgBQKi.js → Commands-DL2kwkRd.js} +1 -1
  60. zenml/zen_server/dashboard/assets/ComponentBadge-D_g62Wv8.js +1 -0
  61. zenml/zen_server/dashboard/assets/{CopyButton-C745BrKi.js → CopyButton-LNcWaa14.js} +1 -1
  62. zenml/zen_server/dashboard/assets/{CsvVizualization-PpAq0CeZ.js → CsvVizualization-DknpE5ej.js} +5 -5
  63. zenml/zen_server/dashboard/assets/{DialogItem-DcVCZEew.js → DialogItem-Bxf8FuAT.js} +1 -1
  64. zenml/zen_server/dashboard/assets/{DisplayDate-BeXgUG_C.js → DisplayDate-CDMUcQHS.js} +1 -1
  65. zenml/zen_server/dashboard/assets/{EmptyState-DeK7H4pr.js → EmptyState-BzdlCwp3.js} +1 -1
  66. zenml/zen_server/dashboard/assets/{Error-BMlzibXj.js → Error-DYflYyps.js} +1 -1
  67. zenml/zen_server/dashboard/assets/ExecutionStatus-C7zyIQKZ.js +1 -0
  68. zenml/zen_server/dashboard/assets/{Helpbox-BLf40fLV.js → Helpbox-oYSGpLqd.js} +1 -1
  69. zenml/zen_server/dashboard/assets/{Infobox-BwisKifi.js → Infobox-Cx4xGoXR.js} +1 -1
  70. zenml/zen_server/dashboard/assets/{InlineAvatar-jEgodSgX.js → InlineAvatar-DiGOWNKF.js} +1 -1
  71. zenml/zen_server/dashboard/assets/{Lock-3lLt1ih0.js → Lock-CYYy18Mm.js} +1 -1
  72. zenml/zen_server/dashboard/assets/{MarkdownVisualization-8O9kTr-2.js → MarkdownVisualization-ylXaAxev.js} +1 -1
  73. zenml/zen_server/dashboard/assets/NestedCollapsible-DYbgyKxK.js +1 -0
  74. zenml/zen_server/dashboard/assets/{NumberBox-T9eELfLZ.js → NumberBox-Dtp3J6g5.js} +1 -1
  75. zenml/zen_server/dashboard/assets/Partials-03iZf8-N.js +1 -0
  76. zenml/zen_server/dashboard/assets/{PasswordChecker-CW0kqY0W.js → PasswordChecker-B0nadgh6.js} +1 -1
  77. zenml/zen_server/dashboard/assets/ProBadge-D_EB8HNo.js +1 -0
  78. zenml/zen_server/dashboard/assets/ProCta-DqNS4v3x.js +1 -0
  79. zenml/zen_server/dashboard/assets/ProviderIcon-Bki2aw8w.js +1 -0
  80. zenml/zen_server/dashboard/assets/{ProviderRadio-BROY1700.js → ProviderRadio-8f43sPD4.js} +1 -1
  81. zenml/zen_server/dashboard/assets/RunSelector-DkPiIiNr.js +1 -0
  82. zenml/zen_server/dashboard/assets/RunsBody-07YEO7qI.js +1 -0
  83. zenml/zen_server/dashboard/assets/SearchField-lp1KgU4e.js +1 -0
  84. zenml/zen_server/dashboard/assets/{SecretTooltip-C_qByGWB.js → SecretTooltip-CgnbyeOx.js} +1 -1
  85. zenml/zen_server/dashboard/assets/{SetPassword-7pRB00El.js → SetPassword-CpP418A2.js} +1 -1
  86. zenml/zen_server/dashboard/assets/StackList-WvuKQusZ.js +1 -0
  87. zenml/zen_server/dashboard/assets/Tabs-BktHkCJJ.js +1 -0
  88. zenml/zen_server/dashboard/assets/Tick-BlMoIlJT.js +1 -0
  89. zenml/zen_server/dashboard/assets/{UpdatePasswordSchemas-DckMEkFf.js → UpdatePasswordSchemas-Sc0A0pP-.js} +1 -1
  90. zenml/zen_server/dashboard/assets/{UsageReason-DVceN14P.js → UsageReason-YYduL4fj.js} +1 -1
  91. zenml/zen_server/dashboard/assets/{WizardFooter-CW0Cvd70.js → WizardFooter-dgmizSJC.js} +1 -1
  92. zenml/zen_server/dashboard/assets/all-pipeline-runs-query-D-c2G6lV.js +1 -0
  93. zenml/zen_server/dashboard/assets/check-DloQpStc.js +1 -0
  94. zenml/zen_server/dashboard/assets/{check-circle-Dwxliy1Z.js → check-circle-jNbX5-sR.js} +1 -1
  95. zenml/zen_server/dashboard/assets/{chevron-down-8wLBS5pQ.js → chevron-down-6JyMkfjR.js} +1 -1
  96. zenml/zen_server/dashboard/assets/{chevron-right-double-DoD8iXWM.js → chevron-right-double-D7ojK9Co.js} +1 -1
  97. zenml/zen_server/dashboard/assets/{code-browser-CZUQs3Wa.js → code-browser-CUFUIHfp.js} +1 -1
  98. zenml/zen_server/dashboard/assets/{copy-CaSMXwiU.js → copy-C8XQA2Ug.js} +1 -1
  99. zenml/zen_server/dashboard/assets/create-stack-DM_JPgef.js +1 -0
  100. zenml/zen_server/dashboard/assets/delete-run-CJdh1P_h.js +1 -0
  101. zenml/zen_server/dashboard/assets/{docker-BFNgg-z3.js → docker-BdA9vrnW.js} +1 -1
  102. zenml/zen_server/dashboard/assets/{dots-horizontal-DK5Duzx4.js → dots-horizontal-otGBOSDJ.js} +1 -1
  103. zenml/zen_server/dashboard/assets/{form-schemas-1AyOCx90.js → form-schemas-K6FYKjwa.js} +1 -1
  104. zenml/zen_server/dashboard/assets/{gcp-7M2Yf3ZK.js → gcp-CFtm4BA7.js} +1 -1
  105. zenml/zen_server/dashboard/assets/{help-Dam461dC.js → help-Cc9bBIJH.js} +1 -1
  106. zenml/zen_server/dashboard/assets/index-B1mVPYxf.js +1 -0
  107. zenml/zen_server/dashboard/assets/index-BAkC7FXi.js +1 -0
  108. zenml/zen_server/dashboard/assets/{index-QQb7wQEC.js → index-CCOPpudF.js} +8 -8
  109. zenml/zen_server/dashboard/assets/index-CEV4Cvaf.js +1 -0
  110. zenml/zen_server/dashboard/assets/index-DlGvJQPn.css +1 -0
  111. zenml/zen_server/dashboard/assets/{index-BVJ8n2-j.js → index-Uu49AX48.js} +1 -1
  112. zenml/zen_server/dashboard/assets/{index.esm-cuVep_NJ.js → index.esm-Dy6Z9Ung.js} +1 -1
  113. zenml/zen_server/dashboard/assets/{kubernetes--g7r02Zu.js → kubernetes-B2wmAJ1d.js} +1 -1
  114. zenml/zen_server/dashboard/assets/{layout-DCSYN7-C.js → layout-BtHBmE4w.js} +1 -1
  115. zenml/zen_server/dashboard/assets/{link-external-CBEk6kEG.js → link-external-b9AXw_sW.js} +1 -1
  116. zenml/zen_server/dashboard/assets/{login-mutation-DTcAFP1l.js → login-mutation-hf-lK87O.js} +1 -1
  117. zenml/zen_server/dashboard/assets/{logs-D5bdJGur.js → logs-WMSM52RF.js} +1 -1
  118. zenml/zen_server/dashboard/assets/{not-found-Cc-JkRH2.js → not-found-BGirLjU-.js} +1 -1
  119. zenml/zen_server/dashboard/assets/{package-Cs35Szwh.js → package-C6uypY4h.js} +1 -1
  120. zenml/zen_server/dashboard/assets/page-0JE_-Ec1.js +1 -0
  121. zenml/zen_server/dashboard/assets/{page-DH_Z7iW1.js → page-6m6yHHlE.js} +1 -1
  122. zenml/zen_server/dashboard/assets/page-BDigxVpo.js +1 -0
  123. zenml/zen_server/dashboard/assets/page-BR68V0V1.js +1 -0
  124. zenml/zen_server/dashboard/assets/page-BRLpxOt0.js +1 -0
  125. zenml/zen_server/dashboard/assets/{page-BQQKaabe.js → page-BU7huvKw.js} +3 -3
  126. zenml/zen_server/dashboard/assets/page-BvqLv2Ky.js +1 -0
  127. zenml/zen_server/dashboard/assets/page-C00YAkaB.js +1 -0
  128. zenml/zen_server/dashboard/assets/{page-N4qoPHKb.js → page-CD-DcWoy.js} +1 -1
  129. zenml/zen_server/dashboard/assets/page-COXXJj1k.js +1 -0
  130. zenml/zen_server/dashboard/assets/page-CbpvrsDL.js +1 -0
  131. zenml/zen_server/dashboard/assets/page-CdMWnQak.js +1 -0
  132. zenml/zen_server/dashboard/assets/{page-ClUVkl-O.js → page-CjGdWY13.js} +1 -1
  133. zenml/zen_server/dashboard/assets/page-CwxrFarU.js +1 -0
  134. zenml/zen_server/dashboard/assets/{page-DLixvR-7.js → page-D01JhjQB.js} +1 -1
  135. zenml/zen_server/dashboard/assets/page-D6uU2ax4.js +1 -0
  136. zenml/zen_server/dashboard/assets/page-D7S3aCbF.js +1 -0
  137. zenml/zen_server/dashboard/assets/{page-9yplj5JT.js → page-DLC-bNBP.js} +1 -1
  138. zenml/zen_server/dashboard/assets/page-DXSTpqRD.js +1 -0
  139. zenml/zen_server/dashboard/assets/{page-DzpVUZ8f.js → page-DakHVWXF.js} +1 -1
  140. zenml/zen_server/dashboard/assets/{page-DIOXwhiD.js → page-Df-Fw0aq.js} +1 -1
  141. zenml/zen_server/dashboard/assets/{page-B-y2XKIc.js → page-DfbXf_8s.js} +1 -1
  142. zenml/zen_server/dashboard/assets/page-DjRJCGb3.js +1 -0
  143. zenml/zen_server/dashboard/assets/{page-C0N5q3l7.js → page-Djikxq_S.js} +1 -1
  144. zenml/zen_server/dashboard/assets/page-Dnovpa0i.js +3 -0
  145. zenml/zen_server/dashboard/assets/page-Dot3LPmL.js +1 -0
  146. zenml/zen_server/dashboard/assets/page-Vcxara9U.js +1 -0
  147. zenml/zen_server/dashboard/assets/page-Xynx4btY.js +14 -0
  148. zenml/zen_server/dashboard/assets/page-YpKAqVSa.js +1 -0
  149. zenml/zen_server/dashboard/assets/page-yYC9OI-E.js +1 -0
  150. zenml/zen_server/dashboard/assets/{persist-DNb5cdrU.js → persist-Coz7ZWvz.js} +1 -1
  151. zenml/zen_server/dashboard/assets/{persist-CP0JmYZ4.js → persist-GjC8PZoC.js} +1 -1
  152. zenml/zen_server/dashboard/assets/{plus-C9IxgN2M.js → plus-tf1V2hTJ.js} +1 -1
  153. zenml/zen_server/dashboard/assets/{refresh-BVu22P_C.js → refresh-BjOeWlEq.js} +1 -1
  154. zenml/zen_server/dashboard/assets/{rocket-CONEmRmB.js → rocket-DjT2cDvG.js} +1 -1
  155. zenml/zen_server/dashboard/assets/sharedSchema-CQb14VSr.js +14 -0
  156. zenml/zen_server/dashboard/assets/stack-detail-query-OPEW-cDJ.js +1 -0
  157. zenml/zen_server/dashboard/assets/{tick-circle-CM1ZScbQ.js → tick-circle-BEX_Tp4v.js} +1 -1
  158. zenml/zen_server/dashboard/assets/{trash-DkJHMOg7.js → trash-arLUMWMS.js} +1 -1
  159. zenml/zen_server/dashboard/assets/{update-server-settings-mutation-DsU8cNVl.js → update-server-settings-mutation-LwuQfHYn.js} +1 -1
  160. zenml/zen_server/dashboard/assets/upgrade-form-CwRHBuXB.webp +0 -0
  161. zenml/zen_server/dashboard/assets/url-CkvKAnwF.js +1 -0
  162. zenml/zen_server/dashboard/assets/{zod-D89GC_vc.js → zod-BwEbpOxH.js} +1 -1
  163. zenml/zen_server/dashboard/index.html +7 -7
  164. zenml/zen_server/deploy/helm/Chart.yaml +1 -1
  165. zenml/zen_server/deploy/helm/README.md +2 -2
  166. zenml/zen_server/exceptions.py +2 -0
  167. zenml/zen_server/jwt.py +30 -13
  168. zenml/zen_server/rbac/endpoint_utils.py +43 -1
  169. zenml/zen_server/rbac/utils.py +0 -2
  170. zenml/zen_server/routers/artifact_version_endpoints.py +27 -1
  171. zenml/zen_server/routers/auth_endpoints.py +134 -102
  172. zenml/zen_server/routers/logs_endpoints.py +66 -0
  173. zenml/zen_server/routers/workspaces_endpoints.py +3 -4
  174. zenml/zen_server/template_execution/utils.py +14 -16
  175. zenml/zen_server/utils.py +27 -0
  176. zenml/zen_server/zen_server_api.py +6 -3
  177. zenml/zen_stores/migrations/versions/0.70.0_release.py +23 -0
  178. zenml/zen_stores/migrations/versions/1cb6477f72d6_move_artifact_save_type.py +99 -0
  179. zenml/zen_stores/migrations/versions/904464ea4041_add_pipeline_model_run_unique_constraints.py +192 -0
  180. zenml/zen_stores/migrations/versions/b557b2871693_update_step_run_input_types.py +33 -0
  181. zenml/zen_stores/rest_zen_store.py +68 -64
  182. zenml/zen_stores/schemas/artifact_schemas.py +8 -1
  183. zenml/zen_stores/schemas/model_schemas.py +27 -3
  184. zenml/zen_stores/schemas/pipeline_run_schemas.py +6 -1
  185. zenml/zen_stores/schemas/pipeline_schemas.py +8 -2
  186. zenml/zen_stores/schemas/run_metadata_schemas.py +1 -48
  187. zenml/zen_stores/schemas/step_run_schemas.py +18 -10
  188. zenml/zen_stores/sql_zen_store.py +283 -219
  189. zenml/zen_stores/zen_store_interface.py +15 -42
  190. {zenml_nightly-0.68.1.dev20241106.dist-info → zenml_nightly-0.70.0.dev20241116.dist-info}/METADATA +2 -2
  191. {zenml_nightly-0.68.1.dev20241106.dist-info → zenml_nightly-0.70.0.dev20241116.dist-info}/RECORD +194 -179
  192. zenml/zen_server/dashboard/assets/AlertDialogDropdownItem-C6N2rGrB.js +0 -1
  193. zenml/zen_server/dashboard/assets/ComponentBadge-DUiEYJHu.js +0 -1
  194. zenml/zen_server/dashboard/assets/ComponentFallbackDialog-BFoH5K4V.js +0 -1
  195. zenml/zen_server/dashboard/assets/ComponentIcon-CAIoUis2.js +0 -1
  196. zenml/zen_server/dashboard/assets/Partials-YPBB3V4q.js +0 -1
  197. zenml/zen_server/dashboard/assets/ProviderIcon-Bb3Xha5A.js +0 -1
  198. zenml/zen_server/dashboard/assets/RunSelector-DCiL3M1c.js +0 -1
  199. zenml/zen_server/dashboard/assets/SearchField-DfUiGFVd.js +0 -1
  200. zenml/zen_server/dashboard/assets/Tick-CykQFPj2.js +0 -1
  201. zenml/zen_server/dashboard/assets/cloud-only-B-s_HMDm.js +0 -1
  202. zenml/zen_server/dashboard/assets/codespaces-BitYDX9d.gif +0 -0
  203. zenml/zen_server/dashboard/assets/create-stack-CEmaPZ4c.js +0 -1
  204. zenml/zen_server/dashboard/assets/delete-run-D-LKbGyz.js +0 -1
  205. zenml/zen_server/dashboard/assets/index-Bpmj40BI.js +0 -1
  206. zenml/zen_server/dashboard/assets/index-CbU4Ln_E.css +0 -1
  207. zenml/zen_server/dashboard/assets/index-DKPhqP2B.js +0 -1
  208. zenml/zen_server/dashboard/assets/page-BBpOxVcY.js +0 -1
  209. zenml/zen_server/dashboard/assets/page-BRInM1Lg.js +0 -1
  210. zenml/zen_server/dashboard/assets/page-BjjlMk7s.js +0 -1
  211. zenml/zen_server/dashboard/assets/page-Bvd7YH2A.js +0 -1
  212. zenml/zen_server/dashboard/assets/page-CT3Nep8W.js +0 -1
  213. zenml/zen_server/dashboard/assets/page-C_f47pBf.js +0 -1
  214. zenml/zen_server/dashboard/assets/page-Cmv8C_yM.js +0 -3
  215. zenml/zen_server/dashboard/assets/page-CyN2bdWG.js +0 -1
  216. zenml/zen_server/dashboard/assets/page-CzzXH4fs.js +0 -1
  217. zenml/zen_server/dashboard/assets/page-DTlGjgnG.js +0 -1
  218. zenml/zen_server/dashboard/assets/page-Dbpl86h0.js +0 -1
  219. zenml/zen_server/dashboard/assets/page-Ddgy6kDS.js +0 -1
  220. zenml/zen_server/dashboard/assets/page-DtCAfBLy.js +0 -9
  221. zenml/zen_server/dashboard/assets/page-Dx16z7nA.js +0 -1
  222. zenml/zen_server/dashboard/assets/page-McUyYbo1.js +0 -1
  223. zenml/zen_server/dashboard/assets/page-T1P3RyAR.js +0 -1
  224. zenml/zen_server/dashboard/assets/page-bKaULTGG.js +0 -1
  225. zenml/zen_server/dashboard/assets/page-sbXUJy9t.js +0 -1
  226. zenml/zen_server/dashboard/assets/sharedSchema-TMLu-nYQ.js +0 -14
  227. zenml/zen_server/dashboard/assets/stack-detail-query-xmYxSsUY.js +0 -1
  228. zenml/zen_server/dashboard/assets/url-D5le3J4q.js +0 -1
  229. zenml/zen_server/routers/run_metadata_endpoints.py +0 -96
  230. {zenml_nightly-0.68.1.dev20241106.dist-info → zenml_nightly-0.70.0.dev20241116.dist-info}/LICENSE +0 -0
  231. {zenml_nightly-0.68.1.dev20241106.dist-info → zenml_nightly-0.70.0.dev20241116.dist-info}/WHEEL +0 -0
  232. {zenml_nightly-0.68.1.dev20241106.dist-info → zenml_nightly-0.70.0.dev20241116.dist-info}/entry_points.txt +0 -0
@@ -13,7 +13,6 @@
13
13
  # permissions and limitations under the License.
14
14
  """Endpoint definitions for authentication (login)."""
15
15
 
16
- from datetime import datetime, timedelta
17
16
  from typing import Optional, Union
18
17
  from urllib.parse import urlencode
19
18
  from uuid import UUID
@@ -40,17 +39,17 @@ from zenml.constants import (
40
39
  VERSION_1,
41
40
  )
42
41
  from zenml.enums import (
42
+ APITokenType,
43
43
  AuthScheme,
44
+ ExecutionStatus,
44
45
  OAuthDeviceStatus,
45
46
  OAuthGrantTypes,
46
47
  )
47
48
  from zenml.exceptions import AuthorizationException
48
49
  from zenml.logger import get_logger
49
50
  from zenml.models import (
50
- APIKeyInternalResponse,
51
51
  OAuthDeviceAuthorizationResponse,
52
52
  OAuthDeviceInternalRequest,
53
- OAuthDeviceInternalResponse,
54
53
  OAuthDeviceInternalUpdate,
55
54
  OAuthDeviceUserAgentHeader,
56
55
  OAuthRedirectResponse,
@@ -63,9 +62,9 @@ from zenml.zen_server.auth import (
63
62
  authenticate_device,
64
63
  authenticate_external_user,
65
64
  authorize,
65
+ generate_access_token,
66
66
  )
67
67
  from zenml.zen_server.exceptions import error_response
68
- from zenml.zen_server.jwt import JWTToken
69
68
  from zenml.zen_server.rate_limit import rate_limit_requests
70
69
  from zenml.zen_server.rbac.models import Action, ResourceType
71
70
  from zenml.zen_server.rbac.utils import verify_permission
@@ -211,68 +210,6 @@ class OAuthLoginRequestForm:
211
210
  )
212
211
 
213
212
 
214
- def generate_access_token(
215
- user_id: UUID,
216
- response: Response,
217
- device: Optional[OAuthDeviceInternalResponse] = None,
218
- api_key: Optional[APIKeyInternalResponse] = None,
219
- ) -> OAuthTokenResponse:
220
- """Generates an access token for the given user.
221
-
222
- Args:
223
- user_id: The ID of the user.
224
- response: The FastAPI response object.
225
- device: The device used for authentication.
226
- api_key: The service account API key used for authentication.
227
-
228
- Returns:
229
- An authentication response with an access token.
230
- """
231
- config = server_config()
232
-
233
- # The JWT tokens are set to expire according to the values configured
234
- # in the server config. Device tokens are handled separately from regular
235
- # user tokens.
236
- expires: Optional[datetime] = None
237
- expires_in: Optional[int] = None
238
- if device:
239
- # If a device was used for authentication, the token will expire
240
- # at the same time as the device.
241
- expires = device.expires
242
- if expires:
243
- expires_in = max(
244
- int(expires.timestamp() - datetime.utcnow().timestamp()), 0
245
- )
246
- elif config.jwt_token_expire_minutes:
247
- expires = datetime.utcnow() + timedelta(
248
- minutes=config.jwt_token_expire_minutes
249
- )
250
- expires_in = config.jwt_token_expire_minutes * 60
251
-
252
- access_token = JWTToken(
253
- user_id=user_id,
254
- device_id=device.id if device else None,
255
- api_key_id=api_key.id if api_key else None,
256
- ).encode(expires=expires)
257
-
258
- if not device:
259
- # Also set the access token as an HTTP only cookie in the response
260
- response.set_cookie(
261
- key=config.get_auth_cookie_name(),
262
- value=access_token,
263
- httponly=True,
264
- samesite="lax",
265
- max_age=config.jwt_token_expire_minutes * 60
266
- if config.jwt_token_expire_minutes
267
- else None,
268
- domain=config.auth_cookie_domain,
269
- )
270
-
271
- return OAuthTokenResponse(
272
- access_token=access_token, expires_in=expires_in, token_type="bearer"
273
- )
274
-
275
-
276
213
  @router.post(
277
214
  LOGIN,
278
215
  response_model=Union[OAuthTokenResponse, OAuthRedirectResponse],
@@ -514,45 +451,74 @@ def device_authorization(
514
451
  )
515
452
  @handle_exceptions
516
453
  def api_token(
517
- pipeline_id: Optional[UUID] = None,
454
+ token_type: APITokenType = APITokenType.GENERIC,
518
455
  schedule_id: Optional[UUID] = None,
519
- expires_minutes: Optional[int] = None,
456
+ pipeline_run_id: Optional[UUID] = None,
457
+ step_run_id: Optional[UUID] = None,
520
458
  auth_context: AuthContext = Security(authorize),
521
459
  ) -> str:
522
- """Get a workload API token for the current user.
460
+ """Generate an API token for the current user.
461
+
462
+ Use this endpoint to generate an API token for the current user. Two types
463
+ of API tokens are supported:
464
+
465
+ * Generic API token: This token is short-lived and can be used for
466
+ generic automation tasks.
467
+ * Workload API token: This token is scoped to a specific pipeline run, step
468
+ run or schedule and is used by pipeline workloads to authenticate with the
469
+ server. A pipeline run ID, step run ID or schedule ID must be provided and
470
+ the generated token will only be valid for the indicated pipeline run, step
471
+ run or schedule. No time limit is imposed on the validity of the token.
472
+ A workload API token can be used to authenticate and generate another
473
+ workload API token, but only for the same schedule, pipeline run ID or step
474
+ run ID, in that order.
523
475
 
524
476
  Args:
525
- pipeline_id: The ID of the pipeline to get the API token for.
526
- schedule_id: The ID of the schedule to get the API token for.
527
- expires_minutes: The number of minutes for which the API token should
528
- be valid. If not provided, the API token will be valid indefinitely.
477
+ token_type: The type of API token to generate.
478
+ schedule_id: The ID of the schedule to scope the workload API token to.
479
+ pipeline_run_id: The ID of the pipeline run to scope the workload API
480
+ token to.
481
+ step_run_id: The ID of the step run to scope the workload API token to.
529
482
  auth_context: The authentication context.
530
483
 
531
484
  Returns:
532
485
  The API token.
533
486
 
534
487
  Raises:
535
- HTTPException: If the user is not authenticated.
536
- AuthorizationException: If trying to scope the API token to a different
537
- pipeline/schedule than the token used to authorize this request.
488
+ AuthorizationException: If not authorized to generate the API token.
489
+ ValueError: If the request is invalid.
538
490
  """
539
491
  token = auth_context.access_token
540
492
  if not token or not auth_context.encoded_access_token:
541
493
  # Should not happen
542
- raise HTTPException(
543
- status_code=status.HTTP_401_UNAUTHORIZED,
544
- detail="Not authenticated.",
545
- )
494
+ raise AuthorizationException("Not authenticated.")
495
+
496
+ if token_type == APITokenType.GENERIC:
497
+ if schedule_id or pipeline_run_id or step_run_id:
498
+ raise ValueError(
499
+ "Generic API tokens cannot be scoped to a schedule, pipeline "
500
+ "run or step run."
501
+ )
502
+
503
+ config = server_config()
504
+
505
+ return generate_access_token(
506
+ user_id=token.user_id,
507
+ expires_in=config.generic_api_token_lifetime,
508
+ ).access_token
546
509
 
547
510
  verify_permission(
548
511
  resource_type=ResourceType.PIPELINE_RUN, action=Action.CREATE
549
512
  )
550
513
 
551
- if pipeline_id and token.pipeline_id and pipeline_id != token.pipeline_id:
552
- raise AuthorizationException(
553
- f"Unable to scope API token to pipeline {pipeline_id}. The "
554
- f"token used to authorize this request is already scoped to "
555
- f"pipeline {token.pipeline_id}."
514
+ schedule_id = schedule_id or token.schedule_id
515
+ pipeline_run_id = pipeline_run_id or token.pipeline_run_id
516
+ step_run_id = step_run_id or token.step_run_id
517
+
518
+ if not pipeline_run_id and not schedule_id and not step_run_id:
519
+ raise ValueError(
520
+ "Workload API tokens must be scoped to a schedule, pipeline run "
521
+ "or step run."
556
522
  )
557
523
 
558
524
  if schedule_id and token.schedule_id and schedule_id != token.schedule_id:
@@ -562,21 +528,87 @@ def api_token(
562
528
  f"schedule {token.schedule_id}."
563
529
  )
564
530
 
565
- if not token.device_id and not token.api_key_id:
566
- # If not authenticated with a device or a service account, the current
567
- # API token is returned as is, without any modifications. Issuing
568
- # workload tokens is only supported for device authenticated users and
569
- # service accounts, because device tokens can be revoked at any time and
570
- # service accounts can be disabled.
571
- return auth_context.encoded_access_token
572
-
573
- # If authenticated with a device, a new API token is generated for the
574
- # pipeline and/or schedule.
575
- if pipeline_id:
576
- token.pipeline_id = pipeline_id
531
+ if (
532
+ pipeline_run_id
533
+ and token.pipeline_run_id
534
+ and pipeline_run_id != token.pipeline_run_id
535
+ ):
536
+ raise AuthorizationException(
537
+ f"Unable to scope API token to pipeline run {pipeline_run_id}. The "
538
+ f"token used to authorize this request is already scoped to "
539
+ f"pipeline run {token.pipeline_run_id}."
540
+ )
541
+
542
+ if step_run_id and token.step_run_id and step_run_id != token.step_run_id:
543
+ raise AuthorizationException(
544
+ f"Unable to scope API token to step run {step_run_id}. The "
545
+ f"token used to authorize this request is already scoped to "
546
+ f"step run {token.step_run_id}."
547
+ )
548
+
577
549
  if schedule_id:
578
- token.schedule_id = schedule_id
579
- expires: Optional[datetime] = None
580
- if expires_minutes:
581
- expires = datetime.utcnow() + timedelta(minutes=expires_minutes)
582
- return token.encode(expires=expires)
550
+ # The schedule must exist
551
+ try:
552
+ schedule = zen_store().get_schedule(schedule_id, hydrate=False)
553
+ except KeyError:
554
+ raise ValueError(
555
+ f"Schedule {schedule_id} does not exist and API tokens cannot "
556
+ "be generated for non-existent schedules for security reasons."
557
+ )
558
+
559
+ if not schedule.active:
560
+ raise ValueError(
561
+ f"Schedule {schedule_id} is not active and API tokens cannot "
562
+ "be generated for inactive schedules for security reasons."
563
+ )
564
+
565
+ if pipeline_run_id:
566
+ # The pipeline run must exist and the run must not be concluded
567
+ try:
568
+ pipeline_run = zen_store().get_run(pipeline_run_id, hydrate=False)
569
+ except KeyError:
570
+ raise ValueError(
571
+ f"Pipeline run {pipeline_run_id} does not exist and API tokens "
572
+ "cannot be generated for non-existent pipeline runs for "
573
+ "security reasons."
574
+ )
575
+
576
+ if pipeline_run.status in [
577
+ ExecutionStatus.FAILED,
578
+ ExecutionStatus.COMPLETED,
579
+ ]:
580
+ raise ValueError(
581
+ f"The execution of pipeline run {pipeline_run_id} has already "
582
+ "concluded and API tokens can no longer be generated for it "
583
+ "for security reasons."
584
+ )
585
+
586
+ if step_run_id:
587
+ # The step run must exist and the step must not be concluded
588
+ try:
589
+ step_run = zen_store().get_run_step(step_run_id, hydrate=False)
590
+ except KeyError:
591
+ raise ValueError(
592
+ f"Step run {step_run_id} does not exist and API tokens cannot "
593
+ "be generated for non-existent step runs for security reasons."
594
+ )
595
+
596
+ if step_run.status in [
597
+ ExecutionStatus.FAILED,
598
+ ExecutionStatus.COMPLETED,
599
+ ]:
600
+ raise ValueError(
601
+ f"The execution of step run {step_run_id} has already "
602
+ "concluded and API tokens can no longer be generated for it "
603
+ "for security reasons."
604
+ )
605
+
606
+ return generate_access_token(
607
+ user_id=token.user_id,
608
+ # Keep the original API key and device token scopes
609
+ api_key=auth_context.api_key,
610
+ device=auth_context.device,
611
+ schedule_id=schedule_id,
612
+ pipeline_run_id=pipeline_run_id,
613
+ step_run_id=step_run_id,
614
+ ).access_token
@@ -0,0 +1,66 @@
1
+ # Copyright (c) ZenML GmbH 2024. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at:
6
+ #
7
+ # https://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
12
+ # or implied. See the License for the specific language governing
13
+ # permissions and limitations under the License.
14
+ """Endpoint definitions for logs."""
15
+
16
+ from uuid import UUID
17
+
18
+ from fastapi import APIRouter, Security
19
+
20
+ from zenml.constants import (
21
+ API,
22
+ LOGS,
23
+ VERSION_1,
24
+ )
25
+ from zenml.models.v2.core.logs import LogsResponse
26
+ from zenml.zen_server.auth import AuthContext, authorize
27
+ from zenml.zen_server.exceptions import error_response
28
+ from zenml.zen_server.rbac.endpoint_utils import (
29
+ verify_permissions_and_get_entity,
30
+ )
31
+ from zenml.zen_server.utils import (
32
+ handle_exceptions,
33
+ zen_store,
34
+ )
35
+
36
+ router = APIRouter(
37
+ prefix=API + VERSION_1 + LOGS,
38
+ tags=["logs"],
39
+ responses={401: error_response, 403: error_response},
40
+ )
41
+
42
+
43
+ @router.get(
44
+ "/{logs_id}",
45
+ response_model=LogsResponse,
46
+ responses={401: error_response, 404: error_response, 422: error_response},
47
+ )
48
+ @handle_exceptions
49
+ def get_logs(
50
+ logs_id: UUID,
51
+ hydrate: bool = True,
52
+ _: AuthContext = Security(authorize),
53
+ ) -> LogsResponse:
54
+ """Returns the requested logs.
55
+
56
+ Args:
57
+ logs_id: ID of the logs.
58
+ hydrate: Flag deciding whether to hydrate the output model(s)
59
+ by including metadata fields in the response.
60
+
61
+ Returns:
62
+ The requested logs.
63
+ """
64
+ return verify_permissions_and_get_entity(
65
+ id=logs_id, get_method=zen_store().get_logs, hydrate=hydrate
66
+ )
@@ -74,7 +74,6 @@ from zenml.models import (
74
74
  PipelineRunRequest,
75
75
  PipelineRunResponse,
76
76
  RunMetadataRequest,
77
- RunMetadataResponse,
78
77
  RunTemplateFilter,
79
78
  RunTemplateRequest,
80
79
  RunTemplateResponse,
@@ -977,7 +976,6 @@ def get_or_create_pipeline_run(
977
976
 
978
977
  @router.post(
979
978
  WORKSPACES + "/{workspace_name_or_id}" + RUN_METADATA,
980
- response_model=List[RunMetadataResponse],
981
979
  responses={401: error_response, 409: error_response, 422: error_response},
982
980
  )
983
981
  @handle_exceptions
@@ -985,7 +983,7 @@ def create_run_metadata(
985
983
  workspace_name_or_id: Union[str, UUID],
986
984
  run_metadata: RunMetadataRequest,
987
985
  auth_context: AuthContext = Security(authorize),
988
- ) -> List[RunMetadataResponse]:
986
+ ) -> None:
989
987
  """Creates run metadata.
990
988
 
991
989
  Args:
@@ -1039,7 +1037,8 @@ def create_run_metadata(
1039
1037
  resource_type=ResourceType.RUN_METADATA, action=Action.CREATE
1040
1038
  )
1041
1039
 
1042
- return zen_store().create_run_metadata(run_metadata)
1040
+ zen_store().create_run_metadata(run_metadata)
1041
+ return None
1043
1042
 
1044
1043
 
1045
1044
  @router.post(
@@ -42,7 +42,7 @@ from zenml.pipelines.run_utils import (
42
42
  )
43
43
  from zenml.stack.flavor import Flavor
44
44
  from zenml.utils import dict_utils, requirements_utils, settings_utils
45
- from zenml.zen_server.auth import AuthContext
45
+ from zenml.zen_server.auth import AuthContext, generate_access_token
46
46
  from zenml.zen_server.template_execution.runner_entrypoint_configuration import (
47
47
  RunnerEntrypointConfiguration,
48
48
  )
@@ -111,17 +111,6 @@ def run_template(
111
111
 
112
112
  new_deployment = zen_store().create_deployment(deployment_request)
113
113
 
114
- if auth_context.access_token:
115
- token = auth_context.access_token
116
- token.pipeline_id = deployment_request.pipeline
117
-
118
- # We create a non-expiring token to make sure its active for the entire
119
- # duration of the pipeline run
120
- api_token = token.encode(expires=None)
121
- else:
122
- assert auth_context.encoded_access_token
123
- api_token = auth_context.encoded_access_token
124
-
125
114
  server_url = server_config().server_url
126
115
  if not server_url:
127
116
  raise RuntimeError(
@@ -130,6 +119,18 @@ def run_template(
130
119
  assert build.zenml_version
131
120
  zenml_version = build.zenml_version
132
121
 
122
+ placeholder_run = create_placeholder_run(deployment=new_deployment)
123
+ assert placeholder_run
124
+
125
+ # We create an API token scoped to the pipeline run
126
+ api_token = generate_access_token(
127
+ user_id=auth_context.user.id,
128
+ pipeline_run_id=placeholder_run.id,
129
+ # Keep the original API key or device scopes, if any
130
+ api_key=auth_context.api_key,
131
+ device=auth_context.device,
132
+ ).access_token
133
+
133
134
  environment = {
134
135
  ENV_ZENML_ACTIVE_WORKSPACE_ID: str(new_deployment.workspace.id),
135
136
  ENV_ZENML_ACTIVE_STACK_ID: str(stack.id),
@@ -145,9 +146,6 @@ def run_template(
145
146
  deployment_id=new_deployment.id
146
147
  )
147
148
 
148
- placeholder_run = create_placeholder_run(deployment=new_deployment)
149
- assert placeholder_run
150
-
151
149
  def _task() -> None:
152
150
  pypi_requirements, apt_packages = (
153
151
  requirements_utils.get_requirements_for_stack(stack=stack)
@@ -350,7 +348,7 @@ def deployment_request_from_template(
350
348
  )
351
349
 
352
350
  step_config_dict_base = pipeline_configuration.model_dump(
353
- exclude={"name", "parameters"}
351
+ exclude={"name", "parameters", "tags"}
354
352
  )
355
353
  steps = {}
356
354
  for invocation_id, step in deployment.step_configurations.items():
zenml/zen_server/utils.py CHANGED
@@ -43,6 +43,7 @@ from zenml.enums import StoreType
43
43
  from zenml.exceptions import IllegalOperationError, OAuthError
44
44
  from zenml.logger import get_logger
45
45
  from zenml.plugins.plugin_flavor_registry import PluginFlavorRegistry
46
+ from zenml.zen_server.cache import MemoryCache
46
47
  from zenml.zen_server.deploy.deployment import (
47
48
  LocalServerDeployment,
48
49
  )
@@ -67,6 +68,7 @@ _rbac: Optional[RBACInterface] = None
67
68
  _feature_gate: Optional[FeatureGateInterface] = None
68
69
  _workload_manager: Optional[WorkloadManagerInterface] = None
69
70
  _plugin_flavor_registry: Optional[PluginFlavorRegistry] = None
71
+ _memcache: Optional[MemoryCache] = None
70
72
 
71
73
 
72
74
  def zen_store() -> "SqlZenStore":
@@ -222,6 +224,31 @@ def initialize_zen_store() -> None:
222
224
  _zen_store = zen_store_
223
225
 
224
226
 
227
+ def initialize_memcache(max_capacity: int, default_expiry: int) -> None:
228
+ """Initialize the memory cache.
229
+
230
+ Args:
231
+ max_capacity: The maximum capacity of the cache.
232
+ default_expiry: The default expiry time in seconds.
233
+ """
234
+ global _memcache
235
+ _memcache = MemoryCache(max_capacity, default_expiry)
236
+
237
+
238
+ def memcache() -> MemoryCache:
239
+ """Return the memory cache.
240
+
241
+ Returns:
242
+ The memory cache.
243
+
244
+ Raises:
245
+ RuntimeError: If the memory cache is not initialized.
246
+ """
247
+ if _memcache is None:
248
+ raise RuntimeError("Memory cache not initialized")
249
+ return _memcache
250
+
251
+
225
252
  _server_config: Optional[ServerConfiguration] = None
226
253
 
227
254
 
@@ -64,13 +64,13 @@ from zenml.zen_server.routers import (
64
64
  devices_endpoints,
65
65
  event_source_endpoints,
66
66
  flavors_endpoints,
67
+ logs_endpoints,
67
68
  model_versions_endpoints,
68
69
  models_endpoints,
69
70
  pipeline_builds_endpoints,
70
71
  pipeline_deployments_endpoints,
71
72
  pipelines_endpoints,
72
73
  plugin_endpoints,
73
- run_metadata_endpoints,
74
74
  run_templates_endpoints,
75
75
  runs_endpoints,
76
76
  schedule_endpoints,
@@ -95,6 +95,7 @@ from zenml.zen_server.secure_headers import (
95
95
  )
96
96
  from zenml.zen_server.utils import (
97
97
  initialize_feature_gate,
98
+ initialize_memcache,
98
99
  initialize_plugins,
99
100
  initialize_rbac,
100
101
  initialize_workload_manager,
@@ -335,9 +336,10 @@ async def infer_source_context(request: Request, call_next: Any) -> Any:
335
336
  @app.on_event("startup")
336
337
  def initialize() -> None:
337
338
  """Initialize the ZenML server."""
339
+ cfg = server_config()
338
340
  # Set the maximum number of worker threads
339
341
  to_thread.current_default_thread_limiter().total_tokens = (
340
- server_config().thread_pool_size
342
+ cfg.thread_pool_size
341
343
  )
342
344
  # IMPORTANT: these need to be run before the fastapi app starts, to avoid
343
345
  # race conditions
@@ -347,6 +349,7 @@ def initialize() -> None:
347
349
  initialize_workload_manager()
348
350
  initialize_plugins()
349
351
  initialize_secure_headers()
352
+ initialize_memcache(cfg.memcache_max_capacity, cfg.memcache_default_expiry)
350
353
 
351
354
 
352
355
  DASHBOARD_REDIRECT_URL = None
@@ -415,6 +418,7 @@ app.include_router(code_repositories_endpoints.router)
415
418
  app.include_router(plugin_endpoints.plugin_router)
416
419
  app.include_router(event_source_endpoints.event_source_router)
417
420
  app.include_router(flavors_endpoints.router)
421
+ app.include_router(logs_endpoints.router)
418
422
  app.include_router(models_endpoints.router)
419
423
  app.include_router(model_versions_endpoints.router)
420
424
  app.include_router(model_versions_endpoints.model_version_artifacts_router)
@@ -423,7 +427,6 @@ app.include_router(pipelines_endpoints.router)
423
427
  app.include_router(pipeline_builds_endpoints.router)
424
428
  app.include_router(pipeline_deployments_endpoints.router)
425
429
  app.include_router(runs_endpoints.router)
426
- app.include_router(run_metadata_endpoints.router)
427
430
  app.include_router(run_templates_endpoints.router)
428
431
  app.include_router(schedule_endpoints.router)
429
432
  app.include_router(secrets_endpoints.router)
@@ -0,0 +1,23 @@
1
+ """Release [0.70.0].
2
+
3
+ Revision ID: 0.70.0
4
+ Revises: 904464ea4041
5
+ Create Date: 2024-11-12 12:20:05.537646
6
+
7
+ """
8
+
9
+ # revision identifiers, used by Alembic.
10
+ revision = "0.70.0"
11
+ down_revision = "904464ea4041"
12
+ branch_labels = None
13
+ depends_on = None
14
+
15
+
16
+ def upgrade() -> None:
17
+ """Upgrade database schema and/or data, creating a new revision."""
18
+ pass
19
+
20
+
21
+ def downgrade() -> None:
22
+ """Downgrade database schema and/or data back to the previous revision."""
23
+ pass