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
@@ -15,7 +15,6 @@
15
15
 
16
16
  import importlib
17
17
  import os
18
- from tempfile import TemporaryDirectory
19
18
  from typing import Any, ClassVar, Dict, Tuple, Type
20
19
 
21
20
  from transformers import (
@@ -46,17 +45,17 @@ class HFPTModelMaterializer(BaseMaterializer):
46
45
  Returns:
47
46
  The model read from the specified dir.
48
47
  """
49
- temp_dir = TemporaryDirectory()
50
- io_utils.copy_dir(
51
- os.path.join(self.uri, DEFAULT_PT_MODEL_DIR), temp_dir.name
52
- )
53
-
54
- config = AutoConfig.from_pretrained(temp_dir.name)
55
- architecture = config.architectures[0]
56
- model_cls = getattr(
57
- importlib.import_module("transformers"), architecture
58
- )
59
- return model_cls.from_pretrained(temp_dir.name)
48
+ with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
49
+ io_utils.copy_dir(
50
+ os.path.join(self.uri, DEFAULT_PT_MODEL_DIR), temp_dir
51
+ )
52
+
53
+ config = AutoConfig.from_pretrained(temp_dir)
54
+ architecture = config.architectures[0]
55
+ model_cls = getattr(
56
+ importlib.import_module("transformers"), architecture
57
+ )
58
+ return model_cls.from_pretrained(temp_dir)
60
59
 
61
60
  def save(self, model: PreTrainedModel) -> None:
62
61
  """Writes a Model to the specified dir.
@@ -64,12 +63,12 @@ class HFPTModelMaterializer(BaseMaterializer):
64
63
  Args:
65
64
  model: The Torch Model to write.
66
65
  """
67
- temp_dir = TemporaryDirectory()
68
- model.save_pretrained(temp_dir.name)
69
- io_utils.copy_dir(
70
- temp_dir.name,
71
- os.path.join(self.uri, DEFAULT_PT_MODEL_DIR),
72
- )
66
+ with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
67
+ model.save_pretrained(temp_dir)
68
+ io_utils.copy_dir(
69
+ temp_dir,
70
+ os.path.join(self.uri, DEFAULT_PT_MODEL_DIR),
71
+ )
73
72
 
74
73
  def extract_metadata(
75
74
  self, model: PreTrainedModel
@@ -14,7 +14,6 @@
14
14
  """Implementation of the Huggingface t5 materializer."""
15
15
 
16
16
  import os
17
- import tempfile
18
17
  from typing import Any, ClassVar, Type, Union
19
18
 
20
19
  from transformers import (
@@ -52,8 +51,7 @@ class HFT5Materializer(BaseMaterializer):
52
51
  ValueError: Unsupported data type used
53
52
  """
54
53
  filepath = self.uri
55
-
56
- with tempfile.TemporaryDirectory(prefix="zenml-temp-") as temp_dir:
54
+ with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
57
55
  # Copy files from artifact store to temporary directory
58
56
  for file in fileio.listdir(filepath):
59
57
  src = os.path.join(filepath, file)
@@ -86,8 +84,7 @@ class HFT5Materializer(BaseMaterializer):
86
84
  Args:
87
85
  obj: A T5ForConditionalGeneration model or T5Tokenizer.
88
86
  """
89
- # Create a temporary directory
90
- with tempfile.TemporaryDirectory(prefix="zenml-temp-") as temp_dir:
87
+ with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
91
88
  # Save the model or tokenizer
92
89
  obj.save_pretrained(temp_dir)
93
90
 
@@ -15,7 +15,6 @@
15
15
 
16
16
  import importlib
17
17
  import os
18
- from tempfile import TemporaryDirectory
19
18
  from typing import Any, ClassVar, Dict, Tuple, Type
20
19
 
21
20
  from transformers import (
@@ -46,17 +45,17 @@ class HFTFModelMaterializer(BaseMaterializer):
46
45
  Returns:
47
46
  The model read from the specified dir.
48
47
  """
49
- temp_dir = TemporaryDirectory()
50
- io_utils.copy_dir(
51
- os.path.join(self.uri, DEFAULT_TF_MODEL_DIR), temp_dir.name
52
- )
53
-
54
- config = AutoConfig.from_pretrained(temp_dir.name)
55
- architecture = "TF" + config.architectures[0]
56
- model_cls = getattr(
57
- importlib.import_module("transformers"), architecture
58
- )
59
- return model_cls.from_pretrained(temp_dir.name)
48
+ with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
49
+ io_utils.copy_dir(
50
+ os.path.join(self.uri, DEFAULT_TF_MODEL_DIR), temp_dir
51
+ )
52
+
53
+ config = AutoConfig.from_pretrained(temp_dir)
54
+ architecture = "TF" + config.architectures[0]
55
+ model_cls = getattr(
56
+ importlib.import_module("transformers"), architecture
57
+ )
58
+ return model_cls.from_pretrained(temp_dir)
60
59
 
61
60
  def save(self, model: TFPreTrainedModel) -> None:
62
61
  """Writes a Model to the specified dir.
@@ -64,12 +63,12 @@ class HFTFModelMaterializer(BaseMaterializer):
64
63
  Args:
65
64
  model: The TF Model to write.
66
65
  """
67
- temp_dir = TemporaryDirectory()
68
- model.save_pretrained(temp_dir.name)
69
- io_utils.copy_dir(
70
- temp_dir.name,
71
- os.path.join(self.uri, DEFAULT_TF_MODEL_DIR),
72
- )
66
+ with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
67
+ model.save_pretrained(temp_dir)
68
+ io_utils.copy_dir(
69
+ temp_dir,
70
+ os.path.join(self.uri, DEFAULT_TF_MODEL_DIR),
71
+ )
73
72
 
74
73
  def extract_metadata(
75
74
  self, model: TFPreTrainedModel
@@ -14,7 +14,6 @@
14
14
  """Implementation of the Huggingface tokenizer materializer."""
15
15
 
16
16
  import os
17
- from tempfile import TemporaryDirectory
18
17
  from typing import Any, ClassVar, Tuple, Type
19
18
 
20
19
  from transformers import AutoTokenizer
@@ -46,7 +45,7 @@ class HFTokenizerMaterializer(BaseMaterializer):
46
45
  Returns:
47
46
  The tokenizer read from the specified dir.
48
47
  """
49
- with TemporaryDirectory() as temp_dir:
48
+ with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
50
49
  io_utils.copy_dir(
51
50
  os.path.join(self.uri, DEFAULT_TOKENIZER_DIR), temp_dir
52
51
  )
@@ -58,7 +57,7 @@ class HFTokenizerMaterializer(BaseMaterializer):
58
57
  Args:
59
58
  tokenizer: The HFTokenizer to write.
60
59
  """
61
- with TemporaryDirectory() as temp_dir:
60
+ with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
62
61
  tokenizer.save_pretrained(temp_dir)
63
62
  io_utils.copy_dir(
64
63
  temp_dir,
@@ -26,7 +26,8 @@ class LangchainIntegration(Integration):
26
26
 
27
27
  NAME = LANGCHAIN
28
28
  REQUIREMENTS = [
29
- "langchain==0.0.325",
29
+ "langchain==0.3.0",
30
+ "langchain-community",
30
31
  "pyyaml>=6.0.1",
31
32
  "tenacity!=8.4.0", # https://github.com/jd/tenacity/issues/471
32
33
  ]
@@ -24,11 +24,37 @@ from zenml.materializers.cloudpickle_materializer import (
24
24
  if TYPE_CHECKING and sys.version_info < (3, 8):
25
25
  OpenAIEmbeddings = Any
26
26
  else:
27
- from langchain.embeddings import OpenAIEmbeddings
27
+ from langchain_community.embeddings import (
28
+ OpenAIEmbeddings,
29
+ )
28
30
 
29
31
 
30
32
  class LangchainOpenaiEmbeddingMaterializer(CloudpickleMaterializer):
31
- """Handle langchain OpenAI embedding objects."""
33
+ """Materializer for Langchain OpenAI Embeddings."""
32
34
 
33
35
  ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.MODEL
34
36
  ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (OpenAIEmbeddings,)
37
+
38
+ def save(self, embeddings: Any) -> None:
39
+ """Saves the embeddings model after clearing non-picklable clients.
40
+
41
+ Args:
42
+ embeddings: The embeddings model to save.
43
+ """
44
+ # Clear the clients which will be recreated on load
45
+ embeddings.client = None
46
+ embeddings.async_client = None
47
+
48
+ # Use the parent class's save implementation which uses cloudpickle
49
+ super().save(embeddings)
50
+
51
+ def load(self, data_type: Type[Any]) -> Any:
52
+ """Loads the embeddings model and lets it recreate clients when needed.
53
+
54
+ Args:
55
+ data_type: The type of the data to load.
56
+
57
+ Returns:
58
+ The loaded embeddings model.
59
+ """
60
+ return super().load(data_type)
@@ -14,7 +14,6 @@
14
14
  """Implementation of the LightGBM booster materializer."""
15
15
 
16
16
  import os
17
- import tempfile
18
17
  from typing import Any, ClassVar, Tuple, Type
19
18
 
20
19
  import lightgbm as lgb
@@ -42,18 +41,13 @@ class LightGBMBoosterMaterializer(BaseMaterializer):
42
41
  A lightgbm Booster object.
43
42
  """
44
43
  filepath = os.path.join(self.uri, DEFAULT_FILENAME)
44
+ with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
45
+ temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
45
46
 
46
- # Create a temporary folder
47
- temp_dir = tempfile.mkdtemp(prefix="zenml-temp-")
48
- temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
49
-
50
- # Copy from artifact store to temporary file
51
- fileio.copy(filepath, temp_file)
52
- booster = lgb.Booster(model_file=temp_file)
53
-
54
- # Cleanup and return
55
- fileio.rmtree(temp_dir)
56
- return booster
47
+ # Copy from artifact store to temporary file
48
+ fileio.copy(filepath, temp_file)
49
+ booster = lgb.Booster(model_file=temp_file)
50
+ return booster
57
51
 
58
52
  def save(self, booster: lgb.Booster) -> None:
59
53
  """Creates a JSON serialization for a lightgbm Booster model.
@@ -62,8 +56,7 @@ class LightGBMBoosterMaterializer(BaseMaterializer):
62
56
  booster: A lightgbm Booster model.
63
57
  """
64
58
  filepath = os.path.join(self.uri, DEFAULT_FILENAME)
65
-
66
- with tempfile.TemporaryDirectory() as tmp_dir:
67
- tmp_path = os.path.join(tmp_dir, "model.txt")
59
+ with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
60
+ tmp_path = os.path.join(temp_dir, "model.txt")
68
61
  booster.save_model(tmp_path)
69
62
  fileio.copy(tmp_path, filepath)
@@ -14,7 +14,6 @@
14
14
  """Implementation of the LightGBM materializer."""
15
15
 
16
16
  import os
17
- import tempfile
18
17
  from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type
19
18
 
20
19
  import lightgbm as lgb
@@ -46,16 +45,14 @@ class LightGBMDatasetMaterializer(BaseMaterializer):
46
45
  """
47
46
  filepath = os.path.join(self.uri, DEFAULT_FILENAME)
48
47
 
49
- # Create a temporary folder
50
- temp_dir = tempfile.mkdtemp(prefix="zenml-temp-")
51
- temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
48
+ with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
49
+ temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
52
50
 
53
- # Copy from artifact store to temporary file
54
- fileio.copy(filepath, temp_file)
55
- matrix = lgb.Dataset(temp_file, free_raw_data=False)
51
+ # Copy from artifact store to temporary file
52
+ fileio.copy(filepath, temp_file)
53
+ matrix = lgb.Dataset(temp_file, free_raw_data=False)
56
54
 
57
- # No clean up this time because matrix is lazy loaded
58
- return matrix
55
+ return matrix
59
56
 
60
57
  def save(self, matrix: lgb.Dataset) -> None:
61
58
  """Creates a binary serialization for a lightgbm.Dataset object.
@@ -65,14 +62,12 @@ class LightGBMDatasetMaterializer(BaseMaterializer):
65
62
  """
66
63
  filepath = os.path.join(self.uri, DEFAULT_FILENAME)
67
64
 
68
- # Make a temporary phantom artifact
69
- temp_dir = tempfile.mkdtemp(prefix="zenml-temp-")
70
- temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
71
- matrix.save_binary(temp_file)
65
+ with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
66
+ temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
67
+ matrix.save_binary(temp_file)
72
68
 
73
- # Copy it into artifact store
74
- fileio.copy(temp_file, filepath)
75
- fileio.rmtree(temp_dir)
69
+ # Copy it into artifact store
70
+ fileio.copy(temp_file, filepath)
76
71
 
77
72
  def extract_metadata(
78
73
  self, matrix: lgb.Dataset
@@ -390,10 +390,13 @@ class LightningOrchestrator(WheeledOrchestrator):
390
390
  studio.run(
391
391
  f"tar -xvzf /teamspace/studios/this_studio/zenml_codes/{filename} -C /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
392
392
  )
393
- studio.upload_file(
394
- env_file_path,
395
- remote_path="/teamspace/studios/this_studio/.lightning_studio/.studiorc",
393
+ studio.upload_file(env_file_path)
394
+ time.sleep(6)
395
+ studio.run(
396
+ f"cp {env_file_path.split('/')[-1]} ./.lightning_studio/.studiorc"
396
397
  )
398
+ studio.run(f"rm {env_file_path.split('/')[-1]}")
399
+
397
400
  studio.run("pip install uv")
398
401
  logger.info(
399
402
  f"Installing requirements: {pipeline_requirements_to_string}"
@@ -481,16 +484,25 @@ class LightningOrchestrator(WheeledOrchestrator):
481
484
  code_path,
482
485
  remote_path=f"/teamspace/studios/this_studio/zenml_codes/{filename}",
483
486
  )
487
+ time.sleep(6)
484
488
  studio.run(
485
489
  f"tar -xvzf /teamspace/studios/this_studio/zenml_codes/{filename} -C /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
486
490
  )
487
491
  logger.info(
488
492
  "Uploading wheel package and installing dependencies on main studio"
489
493
  )
490
- studio.upload_file(
491
- env_file_path,
492
- remote_path="/teamspace/studios/this_studio/.lightning_studio/.studiorc",
494
+ # for some reason uploading file directly to /teamspace/studios/this_studio/.lightning_studio/.studiorc
495
+ # doesn't work, while other file names can be uploaded just fine.
496
+ # below is a workaround to copy the file to the correct location.
497
+ # we first upload it under a different name and then copy it to the
498
+ # right location.
499
+ studio.upload_file(env_file_path)
500
+ time.sleep(6)
501
+ studio.run(
502
+ f"cp {env_file_path.split('/')[-1]} ./.lightning_studio/.studiorc"
493
503
  )
504
+ studio.run(f"rm {env_file_path.split('/')[-1]}")
505
+
494
506
  studio.run("pip install uv")
495
507
  studio.run(f"uv pip install {requirements}")
496
508
  studio.run("pip install zenml")
@@ -557,13 +569,21 @@ class LightningOrchestrator(WheeledOrchestrator):
557
569
  studio.run(
558
570
  f"mkdir -p /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
559
571
  )
560
- studio.upload_file(code_path, remote_path=f"/zenml_codes/{filename}")
572
+ studio.upload_file(
573
+ code_path,
574
+ remote_path=f"/teamspace/studios/this_studio/zenml_codes/{filename}",
575
+ )
576
+ time.sleep(6)
561
577
  studio.run(
562
578
  f"tar -xvzf /teamspace/studios/this_studio/zenml_codes/{filename} -C /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
563
579
  )
564
- studio.upload_file(
565
- env_file_path, remote_path=".lightning_studio/.studiorc"
580
+ studio.upload_file(env_file_path)
581
+ time.sleep(6)
582
+ studio.run(
583
+ f"cp {env_file_path.split('/')[-1]} ./.lightning_studio/.studiorc"
566
584
  )
585
+ studio.run(f"rm {env_file_path.split('/')[-1]}")
586
+
567
587
  studio.run("pip install uv")
568
588
  studio.run(f"uv pip install {details['requirements']}")
569
589
  studio.run("pip install zenml")
@@ -21,7 +21,7 @@ class OpenAIIntegration(Integration):
21
21
  """Definition of OpenAI integration for ZenML."""
22
22
 
23
23
  NAME = OPEN_AI
24
- REQUIREMENTS = ["openai>=0.27.0,<1.0.0"]
24
+ REQUIREMENTS = ["openai>=1.0.0"]
25
25
 
26
26
 
27
27
  OpenAIIntegration.check_installation()
@@ -15,8 +15,9 @@
15
15
 
16
16
  import io
17
17
  import sys
18
+ from typing import Optional
18
19
 
19
- import openai
20
+ from openai import OpenAI
20
21
  from rich.console import Console
21
22
 
22
23
  from zenml import get_step_context
@@ -38,6 +39,8 @@ def openai_alerter_failure_hook_helper(
38
39
  Args:
39
40
  exception: The exception that was raised.
40
41
  model_name: The OpenAI model to use for the chatbot.
42
+
43
+ This implementation uses the OpenAI v1 SDK with automatic retries and backoff.
41
44
  """
42
45
  client = Client()
43
46
  context = get_step_context()
@@ -47,12 +50,15 @@ def openai_alerter_failure_hook_helper(
47
50
  openai_secret = client.get_secret(
48
51
  "openai", allow_partial_name_match=False
49
52
  )
50
- openai_api_key = openai_secret.secret_values.get("api_key")
53
+ openai_api_key: Optional[str] = openai_secret.secret_values.get(
54
+ "api_key"
55
+ )
51
56
  except (KeyError, NotImplementedError):
52
57
  openai_api_key = None
53
58
 
54
59
  alerter = client.active_stack.alerter
55
60
  if alerter and openai_api_key:
61
+ # Capture rich traceback
56
62
  output_captured = io.StringIO()
57
63
  original_stdout = sys.stdout
58
64
  sys.stdout = output_captured
@@ -62,25 +68,44 @@ def openai_alerter_failure_hook_helper(
62
68
  sys.stdout = original_stdout
63
69
  rich_traceback = output_captured.getvalue()
64
70
 
65
- response = openai.ChatCompletion.create( # type: ignore
71
+ # Initialize OpenAI client with timeout and retry settings
72
+ openai_client = OpenAI(
73
+ api_key=openai_api_key,
74
+ max_retries=3, # Will retry 3 times with exponential backoff
75
+ timeout=60.0, # 60 second timeout
76
+ )
77
+
78
+ # Create chat completion using the new client pattern
79
+ response = openai_client.chat.completions.create(
66
80
  model=model_name,
67
81
  messages=[
68
82
  {
69
83
  "role": "user",
70
- "content": f"This is an error message (following an exception of type '{type(exception)}') I encountered while executing a ZenML step. Please suggest ways I might fix the problem. Feel free to give code snippets as examples, and note that your response will be piped to a Slack bot so make sure the formatting is appropriate: {exception} -- {rich_traceback}. Thank you!",
84
+ "content": f"This is an error message (following an exception of type '{type(exception)}') "
85
+ f"I encountered while executing a ZenML step. Please suggest ways I might fix the problem. "
86
+ f"Feel free to give code snippets as examples, and note that your response will be piped "
87
+ f"to a Slack bot so make sure the formatting is appropriate: {exception} -- {rich_traceback}. "
88
+ f"Thank you!",
71
89
  }
72
90
  ],
73
91
  )
74
- suggestion = response["choices"][0]["message"]["content"]
75
- message = "*Failure Hook Notification! Step failed!*" + "\n\n"
76
- message += f"Run name: `{context.pipeline_run.name}`" + "\n"
77
- message += f"Step name: `{context.step_run.name}`" + "\n"
78
- message += f"Parameters: `{context.step_run.config.parameters}`" + "\n"
79
- message += f"Exception: `({type(exception)}) {exception}`" + "\n\n"
80
- message += (
81
- f"*OpenAI ChatGPT's suggestion (model = `{model_name}`) on how to fix it:*\n `{suggestion}`"
82
- + "\n"
92
+
93
+ suggestion = response.choices[0].message.content
94
+
95
+ # Format the alert message
96
+ message = "\n".join(
97
+ [
98
+ "*Failure Hook Notification! Step failed!*",
99
+ "",
100
+ f"Run name: `{context.pipeline_run.name}`",
101
+ f"Step name: `{context.step_run.name}`",
102
+ f"Parameters: `{context.step_run.config.parameters}`",
103
+ f"Exception: `({type(exception)}) {exception}`",
104
+ "",
105
+ f"*OpenAI ChatGPT's suggestion (model = `{model_name}`) on how to fix it:*\n `{suggestion}`",
106
+ ]
83
107
  )
108
+
84
109
  alerter.post(message)
85
110
  elif not openai_api_key:
86
111
  logger.warning(
@@ -111,4 +136,4 @@ def openai_gpt4_alerter_failure_hook(
111
136
  Args:
112
137
  exception: The exception that was raised.
113
138
  """
114
- openai_alerter_failure_hook_helper(exception, "gpt-4")
139
+ openai_alerter_failure_hook_helper(exception, "gpt-4o")
@@ -14,7 +14,6 @@
14
14
  """Materializer for Pillow Image objects."""
15
15
 
16
16
  import os
17
- import tempfile
18
17
  from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type
19
18
 
20
19
  from PIL import Image
@@ -57,16 +56,15 @@ class PillowImageMaterializer(BaseMaterializer):
57
56
  files = io_utils.find_files(self.uri, f"{DEFAULT_IMAGE_FILENAME}.*")
58
57
  filepath = [file for file in files if not fileio.isdir(file)][0]
59
58
 
60
- # create a temporary folder
61
- temp_dir = tempfile.TemporaryDirectory(prefix="zenml-temp-")
62
- temp_file = os.path.join(
63
- temp_dir.name,
64
- f"{DEFAULT_IMAGE_FILENAME}{os.path.splitext(filepath)[1]}",
65
- )
59
+ with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
60
+ temp_file = os.path.join(
61
+ temp_dir,
62
+ f"{DEFAULT_IMAGE_FILENAME}{os.path.splitext(filepath)[1]}",
63
+ )
66
64
 
67
- # copy from artifact store to temporary file
68
- fileio.copy(filepath, temp_file)
69
- return Image.open(temp_file)
65
+ # copy from artifact store to temporary file
66
+ fileio.copy(filepath, temp_file)
67
+ return Image.open(temp_file)
70
68
 
71
69
  def save(self, image: Image.Image) -> None:
72
70
  """Write to artifact store.
@@ -74,18 +72,17 @@ class PillowImageMaterializer(BaseMaterializer):
74
72
  Args:
75
73
  image: An Image.Image object.
76
74
  """
77
- temp_dir = tempfile.TemporaryDirectory(prefix="zenml-temp-")
78
- file_extension = image.format or DEFAULT_IMAGE_EXTENSION
79
- full_filename = f"{DEFAULT_IMAGE_FILENAME}.{file_extension}"
80
- temp_image_path = os.path.join(temp_dir.name, full_filename)
75
+ with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
76
+ file_extension = image.format or DEFAULT_IMAGE_EXTENSION
77
+ full_filename = f"{DEFAULT_IMAGE_FILENAME}.{file_extension}"
78
+ temp_image_path = os.path.join(temp_dir, full_filename)
81
79
 
82
- # save the image in a temporary directory
83
- image.save(temp_image_path)
80
+ # save the image in a temporary directory
81
+ image.save(temp_image_path)
84
82
 
85
- # copy the saved image to the artifact store
86
- artifact_store_path = os.path.join(self.uri, full_filename)
87
- io_utils.copy(temp_image_path, artifact_store_path, overwrite=True) # type: ignore[attr-defined]
88
- temp_dir.cleanup()
83
+ # copy the saved image to the artifact store
84
+ artifact_store_path = os.path.join(self.uri, full_filename)
85
+ io_utils.copy(temp_image_path, artifact_store_path, overwrite=True) # type: ignore[attr-defined]
89
86
 
90
87
  def save_visualizations(
91
88
  self, image: Image.Image
@@ -14,7 +14,6 @@
14
14
  """Polars materializer."""
15
15
 
16
16
  import os
17
- import tempfile
18
17
  from typing import Any, ClassVar, Tuple, Type, Union
19
18
 
20
19
  import polars as pl
@@ -22,7 +21,6 @@ import pyarrow as pa # type: ignore
22
21
  import pyarrow.parquet as pq # type: ignore
23
22
 
24
23
  from zenml.enums import ArtifactType
25
- from zenml.io import fileio
26
24
  from zenml.materializers.base_materializer import BaseMaterializer
27
25
  from zenml.utils import io_utils
28
26
 
@@ -45,35 +43,29 @@ class PolarsMaterializer(BaseMaterializer):
45
43
  Returns:
46
44
  A Polars data frame or series.
47
45
  """
48
- # Create a temporary directory to store the model
49
- temp_dir = tempfile.TemporaryDirectory()
46
+ with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
47
+ io_utils.copy_dir(self.uri, temp_dir)
50
48
 
51
- # Copy from artifact store to temporary directory
52
- io_utils.copy_dir(self.uri, temp_dir.name)
53
-
54
- # Load the data from the temporary directory
55
- table = pq.read_table(
56
- os.path.join(temp_dir.name, "dataframe.parquet").replace("\\", "/")
57
- )
58
-
59
- # If the data is of type pl.Series, convert it back to a pyarrow array
60
- # instead of a table.
61
- if (
62
- table.schema.metadata
63
- and b"zenml_is_pl_series" in table.schema.metadata
64
- ):
65
- isinstance_bytes = table.schema.metadata[b"zenml_is_pl_series"]
66
- isinstance_series = bool.from_bytes(isinstance_bytes, "big")
67
- if isinstance_series:
68
- table = table.column(0)
49
+ # Load the data from the temporary directory
50
+ table = pq.read_table(
51
+ os.path.join(temp_dir, "dataframe.parquet").replace("\\", "/")
52
+ )
69
53
 
70
- # Convert the table to a Polars data frame or series
71
- data = pl.from_arrow(table)
54
+ # If the data is of type pl.Series, convert it back to a pyarrow array
55
+ # instead of a table.
56
+ if (
57
+ table.schema.metadata
58
+ and b"zenml_is_pl_series" in table.schema.metadata
59
+ ):
60
+ isinstance_bytes = table.schema.metadata[b"zenml_is_pl_series"]
61
+ isinstance_series = bool.from_bytes(isinstance_bytes, "big")
62
+ if isinstance_series:
63
+ table = table.column(0)
72
64
 
73
- # Cleanup and return
74
- fileio.rmtree(temp_dir.name)
65
+ # Convert the table to a Polars data frame or series
66
+ data = pl.from_arrow(table)
75
67
 
76
- return data
68
+ return data
77
69
 
78
70
  def save(self, data: Union[pl.DataFrame, pl.Series]) -> None:
79
71
  """Writes Polars data to the artifact store.
@@ -107,15 +99,10 @@ class PolarsMaterializer(BaseMaterializer):
107
99
  {b"zenml_is_pl_series": isinstance_bytes}
108
100
  )
109
101
 
110
- # Create a temporary directory to store the model
111
- temp_dir = tempfile.TemporaryDirectory()
112
-
113
- # Write the table to a Parquet file
114
- path = os.path.join(temp_dir.name, "dataframe.parquet").replace(
115
- "\\", "/"
116
- )
117
- pq.write_table(table, path) # Uses lz4 compression by default
118
- io_utils.copy_dir(temp_dir.name, self.uri)
119
-
120
- # Remove the temporary directory
121
- fileio.rmtree(temp_dir.name)
102
+ with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
103
+ # Write the table to a Parquet file
104
+ path = os.path.join(temp_dir, "dataframe.parquet").replace(
105
+ "\\", "/"
106
+ )
107
+ pq.write_table(table, path) # Uses lz4 compression by default
108
+ io_utils.copy_dir(temp_dir, self.uri)