zenml-nightly 0.68.1.dev20241103__py3-none-any.whl → 0.70.0.dev20241115__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 (264) 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/{load_directory_materializer.py → preexisting_data_materializer.py} +8 -9
  7. zenml/artifacts/utils.py +139 -80
  8. zenml/cli/base.py +4 -4
  9. zenml/cli/model.py +1 -6
  10. zenml/cli/stack.py +1 -0
  11. zenml/client.py +29 -74
  12. zenml/config/server_config.py +17 -1
  13. zenml/constants.py +2 -7
  14. zenml/data_validators/base_data_validator.py +2 -2
  15. zenml/enums.py +20 -4
  16. zenml/exceptions.py +4 -0
  17. zenml/integrations/__init__.py +3 -1
  18. zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py +20 -18
  19. zenml/integrations/azure/orchestrators/azureml_orchestrator.py +1 -1
  20. zenml/integrations/bentoml/materializers/bentoml_bento_materializer.py +19 -31
  21. zenml/integrations/constants.py +1 -0
  22. zenml/integrations/deepchecks/data_validators/deepchecks_data_validator.py +1 -1
  23. zenml/integrations/evidently/__init__.py +1 -1
  24. zenml/integrations/gcp/orchestrators/vertex_orchestrator.py +1 -1
  25. zenml/integrations/huggingface/materializers/huggingface_datasets_materializer.py +8 -12
  26. zenml/integrations/huggingface/materializers/huggingface_pt_model_materializer.py +17 -18
  27. zenml/integrations/huggingface/materializers/huggingface_t5_materializer.py +2 -5
  28. zenml/integrations/huggingface/materializers/huggingface_tf_model_materializer.py +17 -18
  29. zenml/integrations/huggingface/materializers/huggingface_tokenizer_materializer.py +2 -3
  30. zenml/integrations/langchain/__init__.py +2 -1
  31. zenml/integrations/langchain/materializers/openai_embedding_materializer.py +28 -2
  32. zenml/integrations/lightgbm/materializers/lightgbm_booster_materializer.py +8 -15
  33. zenml/integrations/lightgbm/materializers/lightgbm_dataset_materializer.py +11 -16
  34. zenml/integrations/lightning/orchestrators/lightning_orchestrator.py +29 -9
  35. zenml/integrations/openai/__init__.py +1 -1
  36. zenml/integrations/openai/hooks/open_ai_failure_hook.py +39 -14
  37. zenml/integrations/pillow/materializers/pillow_image_materializer.py +17 -20
  38. zenml/integrations/polars/materializers/dataframe_materializer.py +26 -39
  39. zenml/integrations/pycaret/materializers/model_materializer.py +7 -22
  40. zenml/integrations/tensorboard/visualizers/tensorboard_visualizer.py +60 -54
  41. zenml/integrations/tensorflow/materializers/keras_materializer.py +11 -22
  42. zenml/integrations/tensorflow/materializers/tf_dataset_materializer.py +8 -15
  43. zenml/integrations/vllm/__init__.py +50 -0
  44. zenml/integrations/vllm/flavors/__init__.py +21 -0
  45. zenml/integrations/vllm/flavors/vllm_model_deployer_flavor.py +91 -0
  46. zenml/integrations/vllm/model_deployers/__init__.py +19 -0
  47. zenml/integrations/vllm/model_deployers/vllm_model_deployer.py +263 -0
  48. zenml/integrations/vllm/services/__init__.py +19 -0
  49. zenml/integrations/vllm/services/vllm_deployment.py +206 -0
  50. zenml/integrations/whylogs/materializers/whylogs_materializer.py +11 -18
  51. zenml/integrations/xgboost/materializers/xgboost_booster_materializer.py +11 -22
  52. zenml/integrations/xgboost/materializers/xgboost_dmatrix_materializer.py +10 -19
  53. zenml/materializers/base_materializer.py +68 -1
  54. zenml/metadata/lazy_load.py +20 -7
  55. zenml/model/model.py +17 -64
  56. zenml/model/utils.py +5 -0
  57. zenml/models/__init__.py +0 -12
  58. zenml/models/v2/base/filter.py +121 -8
  59. zenml/models/v2/core/artifact_version.py +42 -7
  60. zenml/models/v2/core/model_version.py +26 -5
  61. zenml/models/v2/core/pipeline_run.py +25 -6
  62. zenml/models/v2/core/run_metadata.py +2 -217
  63. zenml/models/v2/core/step_run.py +62 -24
  64. zenml/orchestrators/base_orchestrator.py +12 -1
  65. zenml/orchestrators/input_utils.py +44 -19
  66. zenml/orchestrators/step_launcher.py +4 -3
  67. zenml/orchestrators/step_run_utils.py +19 -15
  68. zenml/orchestrators/step_runner.py +25 -14
  69. zenml/orchestrators/utils.py +45 -26
  70. zenml/stack/flavor.py +9 -5
  71. zenml/stack_deployments/aws_stack_deployment.py +23 -6
  72. zenml/stack_deployments/azure_stack_deployment.py +28 -5
  73. zenml/stack_deployments/gcp_stack_deployment.py +25 -8
  74. zenml/stack_deployments/stack_deployment.py +3 -5
  75. zenml/steps/base_step.py +1 -1
  76. zenml/steps/entrypoint_function_utils.py +3 -5
  77. zenml/steps/step_context.py +5 -2
  78. zenml/steps/utils.py +13 -2
  79. zenml/utils/callback_registry.py +71 -0
  80. zenml/utils/metadata_utils.py +335 -0
  81. zenml/zen_server/auth.py +221 -3
  82. zenml/zen_server/cache.py +208 -0
  83. zenml/zen_server/dashboard/assets/{404-DT4QRUqN.js → 404-NVXKFp-x.js} +1 -1
  84. zenml/zen_server/dashboard/assets/{@radix-DP6vWzyx.js → @radix-DeK6qiuw.js} +1 -1
  85. zenml/zen_server/dashboard/assets/{@react-router-BMhZulnd.js → @react-router-B3Z5rLr2.js} +1 -1
  86. zenml/zen_server/dashboard/assets/{@reactflow-8U9qNlMR.js → @reactflow-CK0KJUen.js} +2 -2
  87. zenml/zen_server/dashboard/assets/{@tanstack-BUCbhJyH.js → @tanstack-DT5WLu9C.js} +1 -1
  88. zenml/zen_server/dashboard/assets/AlertDialogDropdownItem-DezXKmDf.js +1 -0
  89. zenml/zen_server/dashboard/assets/{CodeSnippet-CqybNv0k.js → CodeSnippet-JzR8CEtw.js} +2 -2
  90. zenml/zen_server/dashboard/assets/{CollapsibleCard-0r_8G2Lj.js → CollapsibleCard-DQW_ktMO.js} +1 -1
  91. zenml/zen_server/dashboard/assets/{Commands-BDjgBQKi.js → Commands-DL2kwkRd.js} +1 -1
  92. zenml/zen_server/dashboard/assets/ComponentBadge-D_g62Wv8.js +1 -0
  93. zenml/zen_server/dashboard/assets/{CopyButton-C745BrKi.js → CopyButton-LNcWaa14.js} +1 -1
  94. zenml/zen_server/dashboard/assets/{CsvVizualization-PpAq0CeZ.js → CsvVizualization-DknpE5ej.js} +5 -5
  95. zenml/zen_server/dashboard/assets/{DialogItem-DcVCZEew.js → DialogItem-Bxf8FuAT.js} +1 -1
  96. zenml/zen_server/dashboard/assets/{DisplayDate-BeXgUG_C.js → DisplayDate-CDMUcQHS.js} +1 -1
  97. zenml/zen_server/dashboard/assets/{EmptyState-DeK7H4pr.js → EmptyState-BzdlCwp3.js} +1 -1
  98. zenml/zen_server/dashboard/assets/{Error-BMlzibXj.js → Error-DYflYyps.js} +1 -1
  99. zenml/zen_server/dashboard/assets/ExecutionStatus-C7zyIQKZ.js +1 -0
  100. zenml/zen_server/dashboard/assets/{Helpbox-BLf40fLV.js → Helpbox-oYSGpLqd.js} +1 -1
  101. zenml/zen_server/dashboard/assets/{Infobox-BwisKifi.js → Infobox-Cx4xGoXR.js} +1 -1
  102. zenml/zen_server/dashboard/assets/{InlineAvatar-jEgodSgX.js → InlineAvatar-DiGOWNKF.js} +1 -1
  103. zenml/zen_server/dashboard/assets/{Lock-3lLt1ih0.js → Lock-CYYy18Mm.js} +1 -1
  104. zenml/zen_server/dashboard/assets/{MarkdownVisualization-8O9kTr-2.js → MarkdownVisualization-ylXaAxev.js} +1 -1
  105. zenml/zen_server/dashboard/assets/NestedCollapsible-DYbgyKxK.js +1 -0
  106. zenml/zen_server/dashboard/assets/{NumberBox-T9eELfLZ.js → NumberBox-Dtp3J6g5.js} +1 -1
  107. zenml/zen_server/dashboard/assets/Partials-03iZf8-N.js +1 -0
  108. zenml/zen_server/dashboard/assets/{PasswordChecker-CW0kqY0W.js → PasswordChecker-B0nadgh6.js} +1 -1
  109. zenml/zen_server/dashboard/assets/ProBadge-D_EB8HNo.js +1 -0
  110. zenml/zen_server/dashboard/assets/ProCta-DqNS4v3x.js +1 -0
  111. zenml/zen_server/dashboard/assets/ProviderIcon-Bki2aw8w.js +1 -0
  112. zenml/zen_server/dashboard/assets/{ProviderRadio-BROY1700.js → ProviderRadio-8f43sPD4.js} +1 -1
  113. zenml/zen_server/dashboard/assets/RunSelector-DkPiIiNr.js +1 -0
  114. zenml/zen_server/dashboard/assets/RunsBody-07YEO7qI.js +1 -0
  115. zenml/zen_server/dashboard/assets/SearchField-lp1KgU4e.js +1 -0
  116. zenml/zen_server/dashboard/assets/{SecretTooltip-C_qByGWB.js → SecretTooltip-CgnbyeOx.js} +1 -1
  117. zenml/zen_server/dashboard/assets/{SetPassword-7pRB00El.js → SetPassword-CpP418A2.js} +1 -1
  118. zenml/zen_server/dashboard/assets/StackList-WvuKQusZ.js +1 -0
  119. zenml/zen_server/dashboard/assets/Tabs-BktHkCJJ.js +1 -0
  120. zenml/zen_server/dashboard/assets/Tick-BlMoIlJT.js +1 -0
  121. zenml/zen_server/dashboard/assets/{UpdatePasswordSchemas-DckMEkFf.js → UpdatePasswordSchemas-Sc0A0pP-.js} +1 -1
  122. zenml/zen_server/dashboard/assets/{UsageReason-DVceN14P.js → UsageReason-YYduL4fj.js} +1 -1
  123. zenml/zen_server/dashboard/assets/{WizardFooter-CW0Cvd70.js → WizardFooter-dgmizSJC.js} +1 -1
  124. zenml/zen_server/dashboard/assets/all-pipeline-runs-query-D-c2G6lV.js +1 -0
  125. zenml/zen_server/dashboard/assets/check-DloQpStc.js +1 -0
  126. zenml/zen_server/dashboard/assets/{check-circle-Dwxliy1Z.js → check-circle-jNbX5-sR.js} +1 -1
  127. zenml/zen_server/dashboard/assets/{chevron-down-8wLBS5pQ.js → chevron-down-6JyMkfjR.js} +1 -1
  128. zenml/zen_server/dashboard/assets/{chevron-right-double-DoD8iXWM.js → chevron-right-double-D7ojK9Co.js} +1 -1
  129. zenml/zen_server/dashboard/assets/{code-browser-CZUQs3Wa.js → code-browser-CUFUIHfp.js} +1 -1
  130. zenml/zen_server/dashboard/assets/{copy-CaSMXwiU.js → copy-C8XQA2Ug.js} +1 -1
  131. zenml/zen_server/dashboard/assets/create-stack-DM_JPgef.js +1 -0
  132. zenml/zen_server/dashboard/assets/delete-run-CJdh1P_h.js +1 -0
  133. zenml/zen_server/dashboard/assets/{docker-BFNgg-z3.js → docker-BdA9vrnW.js} +1 -1
  134. zenml/zen_server/dashboard/assets/{dots-horizontal-DK5Duzx4.js → dots-horizontal-otGBOSDJ.js} +1 -1
  135. zenml/zen_server/dashboard/assets/{form-schemas-1AyOCx90.js → form-schemas-K6FYKjwa.js} +1 -1
  136. zenml/zen_server/dashboard/assets/{gcp-7M2Yf3ZK.js → gcp-CFtm4BA7.js} +1 -1
  137. zenml/zen_server/dashboard/assets/{help-Dam461dC.js → help-Cc9bBIJH.js} +1 -1
  138. zenml/zen_server/dashboard/assets/index-B1mVPYxf.js +1 -0
  139. zenml/zen_server/dashboard/assets/index-BAkC7FXi.js +1 -0
  140. zenml/zen_server/dashboard/assets/{index-QQb7wQEC.js → index-CCOPpudF.js} +8 -8
  141. zenml/zen_server/dashboard/assets/index-CEV4Cvaf.js +1 -0
  142. zenml/zen_server/dashboard/assets/index-DlGvJQPn.css +1 -0
  143. zenml/zen_server/dashboard/assets/{index-BVJ8n2-j.js → index-Uu49AX48.js} +1 -1
  144. zenml/zen_server/dashboard/assets/{index.esm-cuVep_NJ.js → index.esm-Dy6Z9Ung.js} +1 -1
  145. zenml/zen_server/dashboard/assets/{kubernetes--g7r02Zu.js → kubernetes-B2wmAJ1d.js} +1 -1
  146. zenml/zen_server/dashboard/assets/{layout-DCSYN7-C.js → layout-BtHBmE4w.js} +1 -1
  147. zenml/zen_server/dashboard/assets/{link-external-CBEk6kEG.js → link-external-b9AXw_sW.js} +1 -1
  148. zenml/zen_server/dashboard/assets/{login-mutation-DTcAFP1l.js → login-mutation-hf-lK87O.js} +1 -1
  149. zenml/zen_server/dashboard/assets/{logs-D5bdJGur.js → logs-WMSM52RF.js} +1 -1
  150. zenml/zen_server/dashboard/assets/{not-found-Cc-JkRH2.js → not-found-BGirLjU-.js} +1 -1
  151. zenml/zen_server/dashboard/assets/{package-Cs35Szwh.js → package-C6uypY4h.js} +1 -1
  152. zenml/zen_server/dashboard/assets/page-0JE_-Ec1.js +1 -0
  153. zenml/zen_server/dashboard/assets/{page-DH_Z7iW1.js → page-6m6yHHlE.js} +1 -1
  154. zenml/zen_server/dashboard/assets/page-BDigxVpo.js +1 -0
  155. zenml/zen_server/dashboard/assets/page-BR68V0V1.js +1 -0
  156. zenml/zen_server/dashboard/assets/page-BRLpxOt0.js +1 -0
  157. zenml/zen_server/dashboard/assets/{page-BQQKaabe.js → page-BU7huvKw.js} +3 -3
  158. zenml/zen_server/dashboard/assets/page-BvqLv2Ky.js +1 -0
  159. zenml/zen_server/dashboard/assets/page-C00YAkaB.js +1 -0
  160. zenml/zen_server/dashboard/assets/{page-N4qoPHKb.js → page-CD-DcWoy.js} +1 -1
  161. zenml/zen_server/dashboard/assets/page-COXXJj1k.js +1 -0
  162. zenml/zen_server/dashboard/assets/page-CbpvrsDL.js +1 -0
  163. zenml/zen_server/dashboard/assets/page-CdMWnQak.js +1 -0
  164. zenml/zen_server/dashboard/assets/{page-ClUVkl-O.js → page-CjGdWY13.js} +1 -1
  165. zenml/zen_server/dashboard/assets/page-CwxrFarU.js +1 -0
  166. zenml/zen_server/dashboard/assets/{page-DLixvR-7.js → page-D01JhjQB.js} +1 -1
  167. zenml/zen_server/dashboard/assets/page-D6uU2ax4.js +1 -0
  168. zenml/zen_server/dashboard/assets/page-D7S3aCbF.js +1 -0
  169. zenml/zen_server/dashboard/assets/{page-9yplj5JT.js → page-DLC-bNBP.js} +1 -1
  170. zenml/zen_server/dashboard/assets/page-DXSTpqRD.js +1 -0
  171. zenml/zen_server/dashboard/assets/{page-DzpVUZ8f.js → page-DakHVWXF.js} +1 -1
  172. zenml/zen_server/dashboard/assets/{page-DIOXwhiD.js → page-Df-Fw0aq.js} +1 -1
  173. zenml/zen_server/dashboard/assets/{page-B-y2XKIc.js → page-DfbXf_8s.js} +1 -1
  174. zenml/zen_server/dashboard/assets/page-DjRJCGb3.js +1 -0
  175. zenml/zen_server/dashboard/assets/{page-C0N5q3l7.js → page-Djikxq_S.js} +1 -1
  176. zenml/zen_server/dashboard/assets/page-Dnovpa0i.js +3 -0
  177. zenml/zen_server/dashboard/assets/page-Dot3LPmL.js +1 -0
  178. zenml/zen_server/dashboard/assets/page-Vcxara9U.js +1 -0
  179. zenml/zen_server/dashboard/assets/page-Xynx4btY.js +14 -0
  180. zenml/zen_server/dashboard/assets/page-YpKAqVSa.js +1 -0
  181. zenml/zen_server/dashboard/assets/page-yYC9OI-E.js +1 -0
  182. zenml/zen_server/dashboard/assets/{persist-DNb5cdrU.js → persist-Coz7ZWvz.js} +1 -1
  183. zenml/zen_server/dashboard/assets/{persist-CP0JmYZ4.js → persist-GjC8PZoC.js} +1 -1
  184. zenml/zen_server/dashboard/assets/{plus-C9IxgN2M.js → plus-tf1V2hTJ.js} +1 -1
  185. zenml/zen_server/dashboard/assets/{refresh-BVu22P_C.js → refresh-BjOeWlEq.js} +1 -1
  186. zenml/zen_server/dashboard/assets/{rocket-CONEmRmB.js → rocket-DjT2cDvG.js} +1 -1
  187. zenml/zen_server/dashboard/assets/sharedSchema-CQb14VSr.js +14 -0
  188. zenml/zen_server/dashboard/assets/stack-detail-query-OPEW-cDJ.js +1 -0
  189. zenml/zen_server/dashboard/assets/{tick-circle-CM1ZScbQ.js → tick-circle-BEX_Tp4v.js} +1 -1
  190. zenml/zen_server/dashboard/assets/{trash-DkJHMOg7.js → trash-arLUMWMS.js} +1 -1
  191. zenml/zen_server/dashboard/assets/{update-server-settings-mutation-DsU8cNVl.js → update-server-settings-mutation-LwuQfHYn.js} +1 -1
  192. zenml/zen_server/dashboard/assets/upgrade-form-CwRHBuXB.webp +0 -0
  193. zenml/zen_server/dashboard/assets/url-CkvKAnwF.js +1 -0
  194. zenml/zen_server/dashboard/assets/{zod-D89GC_vc.js → zod-BwEbpOxH.js} +1 -1
  195. zenml/zen_server/dashboard/index.html +7 -7
  196. zenml/zen_server/deploy/helm/Chart.yaml +1 -1
  197. zenml/zen_server/deploy/helm/README.md +2 -2
  198. zenml/zen_server/exceptions.py +2 -0
  199. zenml/zen_server/jwt.py +30 -13
  200. zenml/zen_server/rbac/endpoint_utils.py +43 -1
  201. zenml/zen_server/rbac/utils.py +0 -2
  202. zenml/zen_server/routers/artifact_version_endpoints.py +27 -1
  203. zenml/zen_server/routers/auth_endpoints.py +134 -102
  204. zenml/zen_server/routers/logs_endpoints.py +66 -0
  205. zenml/zen_server/routers/workspaces_endpoints.py +3 -4
  206. zenml/zen_server/template_execution/utils.py +14 -16
  207. zenml/zen_server/utils.py +27 -0
  208. zenml/zen_server/zen_server_api.py +6 -3
  209. zenml/zen_stores/migrations/versions/0.70.0_release.py +23 -0
  210. zenml/zen_stores/migrations/versions/1cb6477f72d6_move_artifact_save_type.py +99 -0
  211. zenml/zen_stores/migrations/versions/904464ea4041_add_pipeline_model_run_unique_constraints.py +192 -0
  212. zenml/zen_stores/migrations/versions/b557b2871693_update_step_run_input_types.py +33 -0
  213. zenml/zen_stores/rest_zen_store.py +68 -64
  214. zenml/zen_stores/schemas/artifact_schemas.py +8 -1
  215. zenml/zen_stores/schemas/model_schemas.py +27 -3
  216. zenml/zen_stores/schemas/pipeline_run_schemas.py +6 -1
  217. zenml/zen_stores/schemas/pipeline_schemas.py +8 -2
  218. zenml/zen_stores/schemas/run_metadata_schemas.py +1 -48
  219. zenml/zen_stores/schemas/step_run_schemas.py +18 -10
  220. zenml/zen_stores/sql_zen_store.py +283 -219
  221. zenml/zen_stores/zen_store_interface.py +15 -42
  222. {zenml_nightly-0.68.1.dev20241103.dist-info → zenml_nightly-0.70.0.dev20241115.dist-info}/METADATA +2 -2
  223. {zenml_nightly-0.68.1.dev20241103.dist-info → zenml_nightly-0.70.0.dev20241115.dist-info}/RECORD +226 -203
  224. zenml/zen_server/dashboard/assets/AlertDialogDropdownItem-C6N2rGrB.js +0 -1
  225. zenml/zen_server/dashboard/assets/ComponentBadge-DUiEYJHu.js +0 -1
  226. zenml/zen_server/dashboard/assets/ComponentFallbackDialog-BFoH5K4V.js +0 -1
  227. zenml/zen_server/dashboard/assets/ComponentIcon-CAIoUis2.js +0 -1
  228. zenml/zen_server/dashboard/assets/Partials-YPBB3V4q.js +0 -1
  229. zenml/zen_server/dashboard/assets/ProviderIcon-Bb3Xha5A.js +0 -1
  230. zenml/zen_server/dashboard/assets/RunSelector-DCiL3M1c.js +0 -1
  231. zenml/zen_server/dashboard/assets/SearchField-DfUiGFVd.js +0 -1
  232. zenml/zen_server/dashboard/assets/Tick-CykQFPj2.js +0 -1
  233. zenml/zen_server/dashboard/assets/cloud-only-B-s_HMDm.js +0 -1
  234. zenml/zen_server/dashboard/assets/codespaces-BitYDX9d.gif +0 -0
  235. zenml/zen_server/dashboard/assets/create-stack-CEmaPZ4c.js +0 -1
  236. zenml/zen_server/dashboard/assets/delete-run-D-LKbGyz.js +0 -1
  237. zenml/zen_server/dashboard/assets/index-Bpmj40BI.js +0 -1
  238. zenml/zen_server/dashboard/assets/index-CbU4Ln_E.css +0 -1
  239. zenml/zen_server/dashboard/assets/index-DKPhqP2B.js +0 -1
  240. zenml/zen_server/dashboard/assets/page-BBpOxVcY.js +0 -1
  241. zenml/zen_server/dashboard/assets/page-BRInM1Lg.js +0 -1
  242. zenml/zen_server/dashboard/assets/page-BjjlMk7s.js +0 -1
  243. zenml/zen_server/dashboard/assets/page-Bvd7YH2A.js +0 -1
  244. zenml/zen_server/dashboard/assets/page-CT3Nep8W.js +0 -1
  245. zenml/zen_server/dashboard/assets/page-C_f47pBf.js +0 -1
  246. zenml/zen_server/dashboard/assets/page-Cmv8C_yM.js +0 -3
  247. zenml/zen_server/dashboard/assets/page-CyN2bdWG.js +0 -1
  248. zenml/zen_server/dashboard/assets/page-CzzXH4fs.js +0 -1
  249. zenml/zen_server/dashboard/assets/page-DTlGjgnG.js +0 -1
  250. zenml/zen_server/dashboard/assets/page-Dbpl86h0.js +0 -1
  251. zenml/zen_server/dashboard/assets/page-Ddgy6kDS.js +0 -1
  252. zenml/zen_server/dashboard/assets/page-DtCAfBLy.js +0 -9
  253. zenml/zen_server/dashboard/assets/page-Dx16z7nA.js +0 -1
  254. zenml/zen_server/dashboard/assets/page-McUyYbo1.js +0 -1
  255. zenml/zen_server/dashboard/assets/page-T1P3RyAR.js +0 -1
  256. zenml/zen_server/dashboard/assets/page-bKaULTGG.js +0 -1
  257. zenml/zen_server/dashboard/assets/page-sbXUJy9t.js +0 -1
  258. zenml/zen_server/dashboard/assets/sharedSchema-TMLu-nYQ.js +0 -14
  259. zenml/zen_server/dashboard/assets/stack-detail-query-xmYxSsUY.js +0 -1
  260. zenml/zen_server/dashboard/assets/url-D5le3J4q.js +0 -1
  261. zenml/zen_server/routers/run_metadata_endpoints.py +0 -96
  262. {zenml_nightly-0.68.1.dev20241103.dist-info → zenml_nightly-0.70.0.dev20241115.dist-info}/LICENSE +0 -0
  263. {zenml_nightly-0.68.1.dev20241103.dist-info → zenml_nightly-0.70.0.dev20241115.dist-info}/WHEEL +0 -0
  264. {zenml_nightly-0.68.1.dev20241103.dist-info → zenml_nightly-0.70.0.dev20241115.dist-info}/entry_points.txt +0 -0
@@ -18,6 +18,7 @@ import json
18
18
  import logging
19
19
  import math
20
20
  import os
21
+ import random
21
22
  import re
22
23
  import sys
23
24
  import time
@@ -86,6 +87,7 @@ from zenml.config.global_config import GlobalConfiguration
86
87
  from zenml.config.pipeline_run_configuration import PipelineRunConfiguration
87
88
  from zenml.config.secrets_store_config import SecretsStoreConfiguration
88
89
  from zenml.config.server_config import ServerConfiguration
90
+ from zenml.config.step_configurations import StepConfiguration, StepSpec
89
91
  from zenml.config.store_config import StoreConfiguration
90
92
  from zenml.constants import (
91
93
  DEFAULT_PASSWORD,
@@ -118,7 +120,6 @@ from zenml.enums import (
118
120
  StackComponentType,
119
121
  StackDeploymentProvider,
120
122
  StepRunInputArtifactType,
121
- StepRunOutputArtifactType,
122
123
  StoreType,
123
124
  TaggableResourceTypes,
124
125
  )
@@ -126,6 +127,7 @@ from zenml.exceptions import (
126
127
  ActionExistsError,
127
128
  AuthorizationException,
128
129
  BackupSecretsStoreNotConfiguredError,
130
+ EntityCreationError,
129
131
  EntityExistsError,
130
132
  EventSourceExistsError,
131
133
  IllegalOperationError,
@@ -216,9 +218,7 @@ from zenml.models import (
216
218
  PipelineRunResponse,
217
219
  PipelineRunUpdate,
218
220
  PipelineUpdate,
219
- RunMetadataFilter,
220
221
  RunMetadataRequest,
221
- RunMetadataResponse,
222
222
  RunTemplateFilter,
223
223
  RunTemplateRequest,
224
224
  RunTemplateResponse,
@@ -374,6 +374,25 @@ logger = get_logger(__name__)
374
374
  ZENML_SQLITE_DB_FILENAME = "zenml.db"
375
375
 
376
376
 
377
+ def exponential_backoff_with_jitter(
378
+ attempt: int, base_duration: float = 0.05
379
+ ) -> float:
380
+ """Exponential backoff with jitter.
381
+
382
+ Implemented the `Full jitter` algorithm described in
383
+ https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/
384
+
385
+ Args:
386
+ attempt: The backoff attempt.
387
+ base_duration: The backoff base duration.
388
+
389
+ Returns:
390
+ The backoff duration.
391
+ """
392
+ exponential_backoff = base_duration * 1.5**attempt
393
+ return random.uniform(0, exponential_backoff)
394
+
395
+
377
396
  class SQLDatabaseDriver(StrEnum):
378
397
  """SQL database drivers supported by the SQL ZenML store."""
379
398
 
@@ -2784,7 +2803,9 @@ class SqlZenStore(BaseZenStore):
2784
2803
  artifact_version: The artifact version to create.
2785
2804
 
2786
2805
  Raises:
2787
- EntityExistsError: If the artifact version already exists.
2806
+ EntityExistsError: If an artifact version with the same name
2807
+ already exists.
2808
+ EntityCreationError: If the artifact version creation failed.
2788
2809
 
2789
2810
  Returns:
2790
2811
  The created artifact version.
@@ -2825,7 +2846,7 @@ class SqlZenStore(BaseZenStore):
2825
2846
  artifact_version_id = artifact_version_schema.id
2826
2847
  except IntegrityError:
2827
2848
  if remaining_tries == 0:
2828
- raise EntityExistsError(
2849
+ raise EntityCreationError(
2829
2850
  f"Failed to create version for artifact "
2830
2851
  f"{artifact_schema.name}. This is most likely "
2831
2852
  "caused by multiple parallel requests that try "
@@ -2833,12 +2854,14 @@ class SqlZenStore(BaseZenStore):
2833
2854
  "database."
2834
2855
  )
2835
2856
  else:
2836
- # Exponential backoff to account for heavy
2837
- # parallelization
2838
- sleep_duration = 0.05 * 1.5 ** (
2857
+ attempt = (
2839
2858
  MAX_RETRIES_FOR_VERSIONED_ENTITY_CREATION
2840
2859
  - remaining_tries
2841
2860
  )
2861
+ sleep_duration = exponential_backoff_with_jitter(
2862
+ attempt=attempt
2863
+ )
2864
+
2842
2865
  logger.debug(
2843
2866
  "Failed to create artifact version %s "
2844
2867
  "(version %s) due to an integrity error. "
@@ -2915,6 +2938,22 @@ class SqlZenStore(BaseZenStore):
2915
2938
  include_metadata=True, include_resources=True
2916
2939
  )
2917
2940
 
2941
+ def batch_create_artifact_versions(
2942
+ self, artifact_versions: List[ArtifactVersionRequest]
2943
+ ) -> List[ArtifactVersionResponse]:
2944
+ """Creates a batch of artifact versions.
2945
+
2946
+ Args:
2947
+ artifact_versions: The artifact versions to create.
2948
+
2949
+ Returns:
2950
+ The created artifact versions.
2951
+ """
2952
+ return [
2953
+ self.create_artifact_version(artifact_version)
2954
+ for artifact_version in artifact_versions
2955
+ ]
2956
+
2918
2957
  def get_artifact_version(
2919
2958
  self, artifact_version_id: UUID, hydrate: bool = True
2920
2959
  ) -> ArtifactVersionResponse:
@@ -4215,20 +4254,6 @@ class SqlZenStore(BaseZenStore):
4215
4254
  EntityExistsError: If an identical pipeline already exists.
4216
4255
  """
4217
4256
  with Session(self.engine) as session:
4218
- # Check if pipeline with the given name already exists
4219
- existing_pipeline = session.exec(
4220
- select(PipelineSchema)
4221
- .where(PipelineSchema.name == pipeline.name)
4222
- .where(PipelineSchema.workspace_id == pipeline.workspace)
4223
- ).first()
4224
- if existing_pipeline is not None:
4225
- raise EntityExistsError(
4226
- f"Unable to create pipeline in workspace "
4227
- f"'{pipeline.workspace}': A pipeline with this name "
4228
- "already exists."
4229
- )
4230
-
4231
- # Create the pipeline
4232
4257
  new_pipeline = PipelineSchema.from_request(pipeline)
4233
4258
 
4234
4259
  if pipeline.tags:
@@ -4239,7 +4264,14 @@ class SqlZenStore(BaseZenStore):
4239
4264
  )
4240
4265
 
4241
4266
  session.add(new_pipeline)
4242
- session.commit()
4267
+ try:
4268
+ session.commit()
4269
+ except IntegrityError:
4270
+ raise EntityExistsError(
4271
+ f"Unable to create pipeline in workspace "
4272
+ f"'{pipeline.workspace}': A pipeline with the name "
4273
+ f"{pipeline.name} already exists."
4274
+ )
4243
4275
  session.refresh(new_pipeline)
4244
4276
 
4245
4277
  return new_pipeline.to_model(
@@ -5083,6 +5115,26 @@ class SqlZenStore(BaseZenStore):
5083
5115
 
5084
5116
  # ----------------------------- Pipeline runs -----------------------------
5085
5117
 
5118
+ def _pipeline_run_exists(self, workspace_id: UUID, name: str) -> bool:
5119
+ """Check if a pipeline name with a certain name exists.
5120
+
5121
+ Args:
5122
+ workspace_id: The workspace to check.
5123
+ name: The run name.
5124
+
5125
+ Returns:
5126
+ If a pipeline run with the given name exists.
5127
+ """
5128
+ with Session(self.engine) as session:
5129
+ return (
5130
+ session.exec(
5131
+ select(PipelineRunSchema.id)
5132
+ .where(PipelineRunSchema.workspace_id == workspace_id)
5133
+ .where(PipelineRunSchema.name == name)
5134
+ ).first()
5135
+ is not None
5136
+ )
5137
+
5086
5138
  def create_run(
5087
5139
  self, pipeline_run: PipelineRunRequest
5088
5140
  ) -> PipelineRunResponse:
@@ -5098,18 +5150,6 @@ class SqlZenStore(BaseZenStore):
5098
5150
  EntityExistsError: If a run with the same name already exists.
5099
5151
  """
5100
5152
  with Session(self.engine) as session:
5101
- # Check if pipeline run with same name already exists.
5102
- existing_domain_run = session.exec(
5103
- select(PipelineRunSchema).where(
5104
- PipelineRunSchema.name == pipeline_run.name
5105
- )
5106
- ).first()
5107
- if existing_domain_run is not None:
5108
- raise EntityExistsError(
5109
- f"Unable to create pipeline run: A pipeline run with name "
5110
- f"'{pipeline_run.name}' already exists."
5111
- )
5112
-
5113
5153
  # Create the pipeline run
5114
5154
  new_run = PipelineRunSchema.from_request(pipeline_run)
5115
5155
 
@@ -5121,7 +5161,22 @@ class SqlZenStore(BaseZenStore):
5121
5161
  )
5122
5162
 
5123
5163
  session.add(new_run)
5124
- session.commit()
5164
+ try:
5165
+ session.commit()
5166
+ except IntegrityError:
5167
+ if self._pipeline_run_exists(
5168
+ workspace_id=pipeline_run.workspace, name=pipeline_run.name
5169
+ ):
5170
+ raise EntityExistsError(
5171
+ f"Unable to create pipeline run: A pipeline run with "
5172
+ f"name '{pipeline_run.name}' already exists."
5173
+ )
5174
+ else:
5175
+ raise EntityExistsError(
5176
+ "Unable to create pipeline run: A pipeline run with "
5177
+ "the same deployment_id and orchestrator_run_id "
5178
+ "already exists."
5179
+ )
5125
5180
 
5126
5181
  return new_run.to_model(
5127
5182
  include_metadata=True, include_resources=True
@@ -5313,21 +5368,14 @@ class SqlZenStore(BaseZenStore):
5313
5368
  if pre_creation_hook:
5314
5369
  pre_creation_hook()
5315
5370
  return self.create_run(pipeline_run), True
5316
- except (EntityExistsError, IntegrityError) as create_error:
5317
- # Creating the run failed with an
5318
- # - IntegrityError: This happens when we violated a unique
5319
- # constraint, which in turn means a run with the same
5320
- # deployment_id and orchestrator_run_id exists. We now fetch and
5321
- # return that run.
5322
- # - EntityExistsError: This happens when a run with the same name
5323
- # already exists. This could be either a different run (in which
5324
- # case we want to fail) or a run created by a step of the same
5325
- # pipeline run (in which case we want to return it).
5326
- # Note: The IntegrityError might also be raised when other unique
5327
- # constraints get violated. The only other such constraint is the
5328
- # primary key constraint on the run ID, which means we randomly
5329
- # generated an existing UUID. In this case the call below will fail,
5330
- # but the chance of that happening is so low we don't handle it.
5371
+ except EntityExistsError as create_error:
5372
+ # Creating the run failed because
5373
+ # - a run with the same deployment_id and orchestrator_run_id
5374
+ # exists. We now fetch and return that run.
5375
+ # - a run with the same name already exists. This could be either a
5376
+ # different run (in which case we want to fail) or a run created
5377
+ # by a step of the same pipeline run (in which case we want to
5378
+ # return it).
5331
5379
  try:
5332
5380
  return (
5333
5381
  self._get_run_by_orchestrator_run_id(
@@ -5337,18 +5385,11 @@ class SqlZenStore(BaseZenStore):
5337
5385
  False,
5338
5386
  )
5339
5387
  except KeyError:
5340
- if isinstance(create_error, EntityExistsError):
5341
- # There was a run with the same name which does not share
5342
- # the deployment_id and orchestrator_run_id -> We fail with
5343
- # the error that run names must be unique.
5344
- raise create_error from None
5345
-
5346
- # This should never happen as the run creation failed with an
5347
- # IntegrityError which means a run with the deployment_id and
5348
- # orchestrator_run_id exists.
5349
- raise RuntimeError(
5350
- f"Failed to get or create run: {create_error}"
5351
- )
5388
+ # We should only get here if the run creation failed because
5389
+ # of a name conflict. We raise the error that happened during
5390
+ # creation in any case to forward the error message to the
5391
+ # user.
5392
+ raise create_error
5352
5393
 
5353
5394
  def list_runs(
5354
5395
  self,
@@ -5465,9 +5506,7 @@ class SqlZenStore(BaseZenStore):
5465
5506
 
5466
5507
  # ----------------------------- Run Metadata -----------------------------
5467
5508
 
5468
- def create_run_metadata(
5469
- self, run_metadata: RunMetadataRequest
5470
- ) -> List[RunMetadataResponse]:
5509
+ def create_run_metadata(self, run_metadata: RunMetadataRequest) -> None:
5471
5510
  """Creates run metadata.
5472
5511
 
5473
5512
  Args:
@@ -5476,7 +5515,6 @@ class SqlZenStore(BaseZenStore):
5476
5515
  Returns:
5477
5516
  The created run metadata.
5478
5517
  """
5479
- return_value: List[RunMetadataResponse] = []
5480
5518
  with Session(self.engine) as session:
5481
5519
  for key, value in run_metadata.values.items():
5482
5520
  type_ = run_metadata.types[key]
@@ -5492,70 +5530,7 @@ class SqlZenStore(BaseZenStore):
5492
5530
  )
5493
5531
  session.add(run_metadata_schema)
5494
5532
  session.commit()
5495
- return_value.append(
5496
- run_metadata_schema.to_model(
5497
- include_metadata=True, include_resources=True
5498
- )
5499
- )
5500
- return return_value
5501
-
5502
- def get_run_metadata(
5503
- self, run_metadata_id: UUID, hydrate: bool = True
5504
- ) -> RunMetadataResponse:
5505
- """Gets run metadata with the given ID.
5506
-
5507
- Args:
5508
- run_metadata_id: The ID of the run metadata to get.
5509
- hydrate: Flag deciding whether to hydrate the output model(s)
5510
- by including metadata fields in the response.
5511
-
5512
- Returns:
5513
- The run metadata.
5514
-
5515
- Raises:
5516
- KeyError: if the run metadata doesn't exist.
5517
- """
5518
- with Session(self.engine) as session:
5519
- run_metadata = session.exec(
5520
- select(RunMetadataSchema).where(
5521
- RunMetadataSchema.id == run_metadata_id
5522
- )
5523
- ).first()
5524
- if run_metadata is None:
5525
- raise KeyError(
5526
- f"Unable to get run metadata with ID "
5527
- f"{run_metadata_id}: "
5528
- f"No run metadata with this ID found."
5529
- )
5530
- return run_metadata.to_model(
5531
- include_metadata=hydrate, include_resources=True
5532
- )
5533
-
5534
- def list_run_metadata(
5535
- self,
5536
- run_metadata_filter_model: RunMetadataFilter,
5537
- hydrate: bool = False,
5538
- ) -> Page[RunMetadataResponse]:
5539
- """List run metadata.
5540
-
5541
- Args:
5542
- run_metadata_filter_model: All filter parameters including
5543
- pagination params.
5544
- hydrate: Flag deciding whether to hydrate the output model(s)
5545
- by including metadata fields in the response.
5546
-
5547
- Returns:
5548
- The run metadata.
5549
- """
5550
- with Session(self.engine) as session:
5551
- query = select(RunMetadataSchema)
5552
- return self.filter_and_paginate(
5553
- session=session,
5554
- query=query,
5555
- table=RunMetadataSchema,
5556
- filter_model=run_metadata_filter_model,
5557
- hydrate=hydrate,
5558
- )
5533
+ return None
5559
5534
 
5560
5535
  # ----------------------------- Schedules -----------------------------
5561
5536
 
@@ -8173,25 +8148,35 @@ class SqlZenStore(BaseZenStore):
8173
8148
  session=session,
8174
8149
  )
8175
8150
 
8151
+ session.commit()
8152
+ session.refresh(step_schema)
8153
+
8154
+ step_model = step_schema.to_model(include_metadata=True)
8155
+
8176
8156
  # Save input artifact IDs into the database.
8177
8157
  for input_name, artifact_version_id in step_run.inputs.items():
8158
+ input_type = self._get_step_run_input_type(
8159
+ input_name=input_name,
8160
+ step_config=step_model.config,
8161
+ step_spec=step_model.spec,
8162
+ )
8178
8163
  self._set_run_step_input_artifact(
8179
8164
  run_step_id=step_schema.id,
8180
8165
  artifact_version_id=artifact_version_id,
8181
8166
  name=input_name,
8182
- input_type=StepRunInputArtifactType.DEFAULT,
8167
+ input_type=input_type,
8183
8168
  session=session,
8184
8169
  )
8185
8170
 
8186
8171
  # Save output artifact IDs into the database.
8187
- for output_name, artifact_version_id in step_run.outputs.items():
8188
- self._set_run_step_output_artifact(
8189
- step_run_id=step_schema.id,
8190
- artifact_version_id=artifact_version_id,
8191
- name=output_name,
8192
- output_type=StepRunOutputArtifactType.DEFAULT,
8193
- session=session,
8194
- )
8172
+ for output_name, artifact_version_ids in step_run.outputs.items():
8173
+ for artifact_version_id in artifact_version_ids:
8174
+ self._set_run_step_output_artifact(
8175
+ step_run_id=step_schema.id,
8176
+ artifact_version_id=artifact_version_id,
8177
+ name=output_name,
8178
+ session=session,
8179
+ )
8195
8180
 
8196
8181
  if step_run.status != ExecutionStatus.RUNNING:
8197
8182
  self._update_pipeline_run_status(
@@ -8199,6 +8184,7 @@ class SqlZenStore(BaseZenStore):
8199
8184
  )
8200
8185
 
8201
8186
  session.commit()
8187
+ session.refresh(step_schema)
8202
8188
 
8203
8189
  return step_schema.to_model(
8204
8190
  include_metadata=True, include_resources=True
@@ -8291,26 +8277,12 @@ class SqlZenStore(BaseZenStore):
8291
8277
  existing_step_run.update(step_run_update)
8292
8278
  session.add(existing_step_run)
8293
8279
 
8294
- # Update the output artifacts.
8280
+ # Update the artifacts.
8295
8281
  for name, artifact_version_id in step_run_update.outputs.items():
8296
8282
  self._set_run_step_output_artifact(
8297
8283
  step_run_id=step_run_id,
8298
8284
  artifact_version_id=artifact_version_id,
8299
8285
  name=name,
8300
- output_type=StepRunOutputArtifactType.DEFAULT,
8301
- session=session,
8302
- )
8303
-
8304
- # Update saved artifacts
8305
- for (
8306
- artifact_name,
8307
- artifact_version_id,
8308
- ) in step_run_update.saved_artifact_versions.items():
8309
- self._set_run_step_output_artifact(
8310
- step_run_id=step_run_id,
8311
- artifact_version_id=artifact_version_id,
8312
- name=artifact_name,
8313
- output_type=StepRunOutputArtifactType.MANUAL,
8314
8286
  session=session,
8315
8287
  )
8316
8288
 
@@ -8339,6 +8311,34 @@ class SqlZenStore(BaseZenStore):
8339
8311
  include_metadata=True, include_resources=True
8340
8312
  )
8341
8313
 
8314
+ def _get_step_run_input_type(
8315
+ self,
8316
+ input_name: str,
8317
+ step_config: StepConfiguration,
8318
+ step_spec: StepSpec,
8319
+ ) -> StepRunInputArtifactType:
8320
+ """Get the input type of an artifact.
8321
+
8322
+ Args:
8323
+ input_name: The name of the input artifact.
8324
+ step_config: The step config.
8325
+ step_spec: The step spec.
8326
+
8327
+ Returns:
8328
+ The input type of the artifact.
8329
+ """
8330
+ if input_name in step_spec.inputs:
8331
+ return StepRunInputArtifactType.STEP_OUTPUT
8332
+ if input_name in step_config.external_input_artifacts:
8333
+ return StepRunInputArtifactType.EXTERNAL
8334
+ elif (
8335
+ input_name in step_config.model_artifacts_or_metadata
8336
+ or input_name in step_config.client_lazy_loaders
8337
+ ):
8338
+ return StepRunInputArtifactType.LAZY_LOADED
8339
+ else:
8340
+ return StepRunInputArtifactType.MANUAL
8341
+
8342
8342
  @staticmethod
8343
8343
  def _set_run_step_parent_step(
8344
8344
  child_id: UUID, parent_id: UUID, session: Session
@@ -8457,7 +8457,6 @@ class SqlZenStore(BaseZenStore):
8457
8457
  step_run_id: UUID,
8458
8458
  artifact_version_id: UUID,
8459
8459
  name: str,
8460
- output_type: StepRunOutputArtifactType,
8461
8460
  session: Session,
8462
8461
  ) -> None:
8463
8462
  """Sets an artifact as an output of a step run.
@@ -8466,7 +8465,6 @@ class SqlZenStore(BaseZenStore):
8466
8465
  step_run_id: The ID of the step run.
8467
8466
  artifact_version_id: The ID of the artifact version.
8468
8467
  name: The name of the output in the step run.
8469
- output_type: In which way the artifact was saved by the step.
8470
8468
  session: The database session to use.
8471
8469
 
8472
8470
  Raises:
@@ -8510,7 +8508,6 @@ class SqlZenStore(BaseZenStore):
8510
8508
  step_id=step_run_id,
8511
8509
  artifact_id=artifact_version_id,
8512
8510
  name=name,
8513
- type=output_type.value,
8514
8511
  )
8515
8512
  session.add(assignment)
8516
8513
 
@@ -10030,19 +10027,10 @@ class SqlZenStore(BaseZenStore):
10030
10027
  The newly created model.
10031
10028
 
10032
10029
  Raises:
10033
- EntityExistsError: If a workspace with the given name already exists.
10030
+ EntityExistsError: If a model with the given name already exists.
10034
10031
  """
10035
10032
  validate_name(model)
10036
10033
  with Session(self.engine) as session:
10037
- existing_model = session.exec(
10038
- select(ModelSchema).where(ModelSchema.name == model.name)
10039
- ).first()
10040
- if existing_model is not None:
10041
- raise EntityExistsError(
10042
- f"Unable to create model {model.name}: "
10043
- "A model with this name already exists."
10044
- )
10045
-
10046
10034
  model_schema = ModelSchema.from_request(model)
10047
10035
  session.add(model_schema)
10048
10036
 
@@ -10052,7 +10040,14 @@ class SqlZenStore(BaseZenStore):
10052
10040
  resource_id=model_schema.id,
10053
10041
  resource_type=TaggableResourceTypes.MODEL,
10054
10042
  )
10055
- session.commit()
10043
+ try:
10044
+ session.commit()
10045
+ except IntegrityError:
10046
+ raise EntityExistsError(
10047
+ f"Unable to create model {model.name}: "
10048
+ "A model with this name already exists."
10049
+ )
10050
+
10056
10051
  return model_schema.to_model(
10057
10052
  include_metadata=True, include_resources=True
10058
10053
  )
@@ -10188,6 +10183,50 @@ class SqlZenStore(BaseZenStore):
10188
10183
 
10189
10184
  # ----------------------------- Model Versions -----------------------------
10190
10185
 
10186
+ def _get_next_numeric_version_for_model(
10187
+ self, session: Session, model_id: UUID
10188
+ ) -> int:
10189
+ """Get the next numeric version for a model.
10190
+
10191
+ Args:
10192
+ session: DB session.
10193
+ model_id: ID of the model for which to get the next numeric
10194
+ version.
10195
+
10196
+ Returns:
10197
+ The next numeric version.
10198
+ """
10199
+ current_max_version = session.exec(
10200
+ select(func.max(ModelVersionSchema.number)).where(
10201
+ ModelVersionSchema.model_id == model_id
10202
+ )
10203
+ ).first()
10204
+
10205
+ if current_max_version is None:
10206
+ return 1
10207
+ else:
10208
+ return int(current_max_version) + 1
10209
+
10210
+ def _model_version_exists(self, model_id: UUID, version: str) -> bool:
10211
+ """Check if a model version with a certain version exists.
10212
+
10213
+ Args:
10214
+ model_id: The model ID of the version.
10215
+ version: The version name.
10216
+
10217
+ Returns:
10218
+ If a model version with the given version name exists.
10219
+ """
10220
+ with Session(self.engine) as session:
10221
+ return (
10222
+ session.exec(
10223
+ select(ModelVersionSchema.id)
10224
+ .where(ModelVersionSchema.model_id == model_id)
10225
+ .where(ModelVersionSchema.name == version)
10226
+ ).first()
10227
+ is not None
10228
+ )
10229
+
10191
10230
  @track_decorator(AnalyticsEvent.CREATED_MODEL_VERSION)
10192
10231
  def create_model_version(
10193
10232
  self, model_version: ModelVersionRequest
@@ -10202,70 +10241,95 @@ class SqlZenStore(BaseZenStore):
10202
10241
 
10203
10242
  Raises:
10204
10243
  ValueError: If `number` is not None during model version creation.
10205
- EntityExistsError: If a workspace with the given name already exists.
10244
+ EntityExistsError: If a model version with the given name already
10245
+ exists.
10246
+ EntityCreationError: If the model version creation failed.
10206
10247
  """
10207
10248
  if model_version.number is not None:
10208
10249
  raise ValueError(
10209
10250
  "`number` field must be None during model version creation."
10210
10251
  )
10211
- with Session(self.engine) as session:
10212
- model_version_ = model_version.model_copy()
10213
- model = self.get_model(model_version_.model)
10214
10252
 
10215
- def _check(tolerance: int = 0) -> None:
10216
- query = session.exec(
10217
- select(ModelVersionSchema)
10218
- .where(ModelVersionSchema.model_id == model.id)
10219
- .where(ModelVersionSchema.name == model_version_.name)
10220
- )
10221
- existing_model_version = query.fetchmany(tolerance + 1)
10222
- if (
10223
- existing_model_version is not None
10224
- and len(existing_model_version) > tolerance
10225
- ):
10226
- raise EntityExistsError(
10227
- f"Unable to create model version {model_version_.name}: "
10228
- f"A model version with this name already exists in {model.name} model."
10229
- )
10253
+ model = self.get_model(model_version.model)
10230
10254
 
10231
- _check()
10232
- all_versions = session.exec(
10233
- select(ModelVersionSchema)
10234
- .where(ModelVersionSchema.model_id == model.id)
10235
- .order_by(ModelVersionSchema.number.desc()) # type: ignore[attr-defined]
10236
- ).first()
10255
+ has_custom_name = model_version.name is not None
10256
+ if has_custom_name:
10257
+ validate_name(model_version)
10237
10258
 
10238
- model_version_.number = (
10239
- all_versions.number + 1 if all_versions else 1
10240
- )
10259
+ model_version_id = None
10241
10260
 
10242
- if model_version_.name is None:
10243
- model_version_.name = str(model_version_.number)
10244
- else:
10245
- validate_name(model_version_)
10261
+ remaining_tries = MAX_RETRIES_FOR_VERSIONED_ENTITY_CREATION
10262
+ while remaining_tries > 0:
10263
+ remaining_tries -= 1
10264
+ try:
10265
+ with Session(self.engine) as session:
10266
+ model_version.number = (
10267
+ self._get_next_numeric_version_for_model(
10268
+ session=session,
10269
+ model_id=model.id,
10270
+ )
10271
+ )
10272
+ if not has_custom_name:
10273
+ model_version.name = str(model_version.number)
10246
10274
 
10247
- model_version_schema = ModelVersionSchema.from_request(
10248
- model_version_
10249
- )
10250
- session.add(model_version_schema)
10275
+ model_version_schema = ModelVersionSchema.from_request(
10276
+ model_version
10277
+ )
10278
+ session.add(model_version_schema)
10279
+ session.commit()
10251
10280
 
10252
- if model_version_.tags:
10253
- self._attach_tags_to_resource(
10254
- tag_names=model_version_.tags,
10255
- resource_id=model_version_schema.id,
10256
- resource_type=TaggableResourceTypes.MODEL_VERSION,
10257
- )
10258
- try:
10259
- _check(1)
10260
- session.commit()
10261
- except EntityExistsError as e:
10262
- session.rollback()
10263
- raise e
10281
+ model_version_id = model_version_schema.id
10282
+ break
10283
+ except IntegrityError:
10284
+ if has_custom_name and self._model_version_exists(
10285
+ model_id=model.id, version=cast(str, model_version.name)
10286
+ ):
10287
+ # We failed not because of a version number conflict,
10288
+ # but because the user requested a version name that
10289
+ # is already taken -> We don't retry anymore but fail
10290
+ # immediately.
10291
+ raise EntityExistsError(
10292
+ f"Unable to create model version "
10293
+ f"{model.name} (version "
10294
+ f"{model_version.name}): A model with the "
10295
+ "same name and version already exists."
10296
+ )
10297
+ elif remaining_tries == 0:
10298
+ raise EntityCreationError(
10299
+ f"Failed to create version for model "
10300
+ f"{model.name}. This is most likely "
10301
+ "caused by multiple parallel requests that try "
10302
+ "to create versions for this model in the "
10303
+ "database."
10304
+ )
10305
+ else:
10306
+ attempt = (
10307
+ MAX_RETRIES_FOR_VERSIONED_ENTITY_CREATION
10308
+ - remaining_tries
10309
+ )
10310
+ sleep_duration = exponential_backoff_with_jitter(
10311
+ attempt=attempt
10312
+ )
10313
+ logger.debug(
10314
+ "Failed to create model version %s "
10315
+ "(version %s) due to an integrity error. "
10316
+ "Retrying in %f seconds.",
10317
+ model.name,
10318
+ model_version.number,
10319
+ sleep_duration,
10320
+ )
10321
+ time.sleep(sleep_duration)
10264
10322
 
10265
- return model_version_schema.to_model(
10266
- include_metadata=True, include_resources=True
10323
+ assert model_version_id
10324
+ if model_version.tags:
10325
+ self._attach_tags_to_resource(
10326
+ tag_names=model_version.tags,
10327
+ resource_id=model_version_id,
10328
+ resource_type=TaggableResourceTypes.MODEL_VERSION,
10267
10329
  )
10268
10330
 
10331
+ return self.get_model_version(model_version_id)
10332
+
10269
10333
  def get_model_version(
10270
10334
  self, model_version_id: UUID, hydrate: bool = True
10271
10335
  ) -> ModelVersionResponse: