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.
- README.md +1 -1
- RELEASE_NOTES.md +77 -0
- zenml/VERSION +1 -1
- zenml/__init__.py +2 -0
- zenml/artifacts/external_artifact.py +2 -1
- zenml/artifacts/{load_directory_materializer.py → preexisting_data_materializer.py} +8 -9
- zenml/artifacts/utils.py +139 -80
- zenml/cli/base.py +4 -4
- zenml/cli/model.py +1 -6
- zenml/cli/stack.py +1 -0
- zenml/client.py +29 -74
- zenml/config/server_config.py +17 -1
- zenml/constants.py +2 -7
- zenml/data_validators/base_data_validator.py +2 -2
- zenml/enums.py +20 -4
- zenml/exceptions.py +4 -0
- zenml/integrations/__init__.py +3 -1
- zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py +20 -18
- zenml/integrations/azure/orchestrators/azureml_orchestrator.py +1 -1
- zenml/integrations/bentoml/materializers/bentoml_bento_materializer.py +19 -31
- zenml/integrations/constants.py +1 -0
- zenml/integrations/deepchecks/data_validators/deepchecks_data_validator.py +1 -1
- zenml/integrations/evidently/__init__.py +1 -1
- zenml/integrations/gcp/orchestrators/vertex_orchestrator.py +1 -1
- zenml/integrations/huggingface/materializers/huggingface_datasets_materializer.py +8 -12
- zenml/integrations/huggingface/materializers/huggingface_pt_model_materializer.py +17 -18
- zenml/integrations/huggingface/materializers/huggingface_t5_materializer.py +2 -5
- zenml/integrations/huggingface/materializers/huggingface_tf_model_materializer.py +17 -18
- zenml/integrations/huggingface/materializers/huggingface_tokenizer_materializer.py +2 -3
- zenml/integrations/langchain/__init__.py +2 -1
- zenml/integrations/langchain/materializers/openai_embedding_materializer.py +28 -2
- zenml/integrations/lightgbm/materializers/lightgbm_booster_materializer.py +8 -15
- zenml/integrations/lightgbm/materializers/lightgbm_dataset_materializer.py +11 -16
- zenml/integrations/lightning/orchestrators/lightning_orchestrator.py +29 -9
- zenml/integrations/openai/__init__.py +1 -1
- zenml/integrations/openai/hooks/open_ai_failure_hook.py +39 -14
- zenml/integrations/pillow/materializers/pillow_image_materializer.py +17 -20
- zenml/integrations/polars/materializers/dataframe_materializer.py +26 -39
- zenml/integrations/pycaret/materializers/model_materializer.py +7 -22
- zenml/integrations/tensorboard/visualizers/tensorboard_visualizer.py +60 -54
- zenml/integrations/tensorflow/materializers/keras_materializer.py +11 -22
- zenml/integrations/tensorflow/materializers/tf_dataset_materializer.py +8 -15
- zenml/integrations/vllm/__init__.py +50 -0
- zenml/integrations/vllm/flavors/__init__.py +21 -0
- zenml/integrations/vllm/flavors/vllm_model_deployer_flavor.py +91 -0
- zenml/integrations/vllm/model_deployers/__init__.py +19 -0
- zenml/integrations/vllm/model_deployers/vllm_model_deployer.py +263 -0
- zenml/integrations/vllm/services/__init__.py +19 -0
- zenml/integrations/vllm/services/vllm_deployment.py +206 -0
- zenml/integrations/whylogs/materializers/whylogs_materializer.py +11 -18
- zenml/integrations/xgboost/materializers/xgboost_booster_materializer.py +11 -22
- zenml/integrations/xgboost/materializers/xgboost_dmatrix_materializer.py +10 -19
- zenml/materializers/base_materializer.py +68 -1
- zenml/metadata/lazy_load.py +20 -7
- zenml/model/model.py +17 -64
- zenml/model/utils.py +5 -0
- zenml/models/__init__.py +0 -12
- zenml/models/v2/base/filter.py +121 -8
- zenml/models/v2/core/artifact_version.py +42 -7
- zenml/models/v2/core/model_version.py +26 -5
- zenml/models/v2/core/pipeline_run.py +25 -6
- zenml/models/v2/core/run_metadata.py +2 -217
- zenml/models/v2/core/step_run.py +62 -24
- zenml/orchestrators/base_orchestrator.py +12 -1
- zenml/orchestrators/input_utils.py +44 -19
- zenml/orchestrators/step_launcher.py +4 -3
- zenml/orchestrators/step_run_utils.py +19 -15
- zenml/orchestrators/step_runner.py +25 -14
- zenml/orchestrators/utils.py +45 -26
- zenml/stack/flavor.py +9 -5
- zenml/stack_deployments/aws_stack_deployment.py +23 -6
- zenml/stack_deployments/azure_stack_deployment.py +28 -5
- zenml/stack_deployments/gcp_stack_deployment.py +25 -8
- zenml/stack_deployments/stack_deployment.py +3 -5
- zenml/steps/base_step.py +1 -1
- zenml/steps/entrypoint_function_utils.py +3 -5
- zenml/steps/step_context.py +5 -2
- zenml/steps/utils.py +13 -2
- zenml/utils/callback_registry.py +71 -0
- zenml/utils/metadata_utils.py +335 -0
- zenml/zen_server/auth.py +221 -3
- zenml/zen_server/cache.py +208 -0
- zenml/zen_server/dashboard/assets/{404-DT4QRUqN.js → 404-NVXKFp-x.js} +1 -1
- zenml/zen_server/dashboard/assets/{@radix-DP6vWzyx.js → @radix-DeK6qiuw.js} +1 -1
- zenml/zen_server/dashboard/assets/{@react-router-BMhZulnd.js → @react-router-B3Z5rLr2.js} +1 -1
- zenml/zen_server/dashboard/assets/{@reactflow-8U9qNlMR.js → @reactflow-CK0KJUen.js} +2 -2
- zenml/zen_server/dashboard/assets/{@tanstack-BUCbhJyH.js → @tanstack-DT5WLu9C.js} +1 -1
- zenml/zen_server/dashboard/assets/AlertDialogDropdownItem-DezXKmDf.js +1 -0
- zenml/zen_server/dashboard/assets/{CodeSnippet-CqybNv0k.js → CodeSnippet-JzR8CEtw.js} +2 -2
- zenml/zen_server/dashboard/assets/{CollapsibleCard-0r_8G2Lj.js → CollapsibleCard-DQW_ktMO.js} +1 -1
- zenml/zen_server/dashboard/assets/{Commands-BDjgBQKi.js → Commands-DL2kwkRd.js} +1 -1
- zenml/zen_server/dashboard/assets/ComponentBadge-D_g62Wv8.js +1 -0
- zenml/zen_server/dashboard/assets/{CopyButton-C745BrKi.js → CopyButton-LNcWaa14.js} +1 -1
- zenml/zen_server/dashboard/assets/{CsvVizualization-PpAq0CeZ.js → CsvVizualization-DknpE5ej.js} +5 -5
- zenml/zen_server/dashboard/assets/{DialogItem-DcVCZEew.js → DialogItem-Bxf8FuAT.js} +1 -1
- zenml/zen_server/dashboard/assets/{DisplayDate-BeXgUG_C.js → DisplayDate-CDMUcQHS.js} +1 -1
- zenml/zen_server/dashboard/assets/{EmptyState-DeK7H4pr.js → EmptyState-BzdlCwp3.js} +1 -1
- zenml/zen_server/dashboard/assets/{Error-BMlzibXj.js → Error-DYflYyps.js} +1 -1
- zenml/zen_server/dashboard/assets/ExecutionStatus-C7zyIQKZ.js +1 -0
- zenml/zen_server/dashboard/assets/{Helpbox-BLf40fLV.js → Helpbox-oYSGpLqd.js} +1 -1
- zenml/zen_server/dashboard/assets/{Infobox-BwisKifi.js → Infobox-Cx4xGoXR.js} +1 -1
- zenml/zen_server/dashboard/assets/{InlineAvatar-jEgodSgX.js → InlineAvatar-DiGOWNKF.js} +1 -1
- zenml/zen_server/dashboard/assets/{Lock-3lLt1ih0.js → Lock-CYYy18Mm.js} +1 -1
- zenml/zen_server/dashboard/assets/{MarkdownVisualization-8O9kTr-2.js → MarkdownVisualization-ylXaAxev.js} +1 -1
- zenml/zen_server/dashboard/assets/NestedCollapsible-DYbgyKxK.js +1 -0
- zenml/zen_server/dashboard/assets/{NumberBox-T9eELfLZ.js → NumberBox-Dtp3J6g5.js} +1 -1
- zenml/zen_server/dashboard/assets/Partials-03iZf8-N.js +1 -0
- zenml/zen_server/dashboard/assets/{PasswordChecker-CW0kqY0W.js → PasswordChecker-B0nadgh6.js} +1 -1
- zenml/zen_server/dashboard/assets/ProBadge-D_EB8HNo.js +1 -0
- zenml/zen_server/dashboard/assets/ProCta-DqNS4v3x.js +1 -0
- zenml/zen_server/dashboard/assets/ProviderIcon-Bki2aw8w.js +1 -0
- zenml/zen_server/dashboard/assets/{ProviderRadio-BROY1700.js → ProviderRadio-8f43sPD4.js} +1 -1
- zenml/zen_server/dashboard/assets/RunSelector-DkPiIiNr.js +1 -0
- zenml/zen_server/dashboard/assets/RunsBody-07YEO7qI.js +1 -0
- zenml/zen_server/dashboard/assets/SearchField-lp1KgU4e.js +1 -0
- zenml/zen_server/dashboard/assets/{SecretTooltip-C_qByGWB.js → SecretTooltip-CgnbyeOx.js} +1 -1
- zenml/zen_server/dashboard/assets/{SetPassword-7pRB00El.js → SetPassword-CpP418A2.js} +1 -1
- zenml/zen_server/dashboard/assets/StackList-WvuKQusZ.js +1 -0
- zenml/zen_server/dashboard/assets/Tabs-BktHkCJJ.js +1 -0
- zenml/zen_server/dashboard/assets/Tick-BlMoIlJT.js +1 -0
- zenml/zen_server/dashboard/assets/{UpdatePasswordSchemas-DckMEkFf.js → UpdatePasswordSchemas-Sc0A0pP-.js} +1 -1
- zenml/zen_server/dashboard/assets/{UsageReason-DVceN14P.js → UsageReason-YYduL4fj.js} +1 -1
- zenml/zen_server/dashboard/assets/{WizardFooter-CW0Cvd70.js → WizardFooter-dgmizSJC.js} +1 -1
- zenml/zen_server/dashboard/assets/all-pipeline-runs-query-D-c2G6lV.js +1 -0
- zenml/zen_server/dashboard/assets/check-DloQpStc.js +1 -0
- zenml/zen_server/dashboard/assets/{check-circle-Dwxliy1Z.js → check-circle-jNbX5-sR.js} +1 -1
- zenml/zen_server/dashboard/assets/{chevron-down-8wLBS5pQ.js → chevron-down-6JyMkfjR.js} +1 -1
- zenml/zen_server/dashboard/assets/{chevron-right-double-DoD8iXWM.js → chevron-right-double-D7ojK9Co.js} +1 -1
- zenml/zen_server/dashboard/assets/{code-browser-CZUQs3Wa.js → code-browser-CUFUIHfp.js} +1 -1
- zenml/zen_server/dashboard/assets/{copy-CaSMXwiU.js → copy-C8XQA2Ug.js} +1 -1
- zenml/zen_server/dashboard/assets/create-stack-DM_JPgef.js +1 -0
- zenml/zen_server/dashboard/assets/delete-run-CJdh1P_h.js +1 -0
- zenml/zen_server/dashboard/assets/{docker-BFNgg-z3.js → docker-BdA9vrnW.js} +1 -1
- zenml/zen_server/dashboard/assets/{dots-horizontal-DK5Duzx4.js → dots-horizontal-otGBOSDJ.js} +1 -1
- zenml/zen_server/dashboard/assets/{form-schemas-1AyOCx90.js → form-schemas-K6FYKjwa.js} +1 -1
- zenml/zen_server/dashboard/assets/{gcp-7M2Yf3ZK.js → gcp-CFtm4BA7.js} +1 -1
- zenml/zen_server/dashboard/assets/{help-Dam461dC.js → help-Cc9bBIJH.js} +1 -1
- zenml/zen_server/dashboard/assets/index-B1mVPYxf.js +1 -0
- zenml/zen_server/dashboard/assets/index-BAkC7FXi.js +1 -0
- zenml/zen_server/dashboard/assets/{index-QQb7wQEC.js → index-CCOPpudF.js} +8 -8
- zenml/zen_server/dashboard/assets/index-CEV4Cvaf.js +1 -0
- zenml/zen_server/dashboard/assets/index-DlGvJQPn.css +1 -0
- zenml/zen_server/dashboard/assets/{index-BVJ8n2-j.js → index-Uu49AX48.js} +1 -1
- zenml/zen_server/dashboard/assets/{index.esm-cuVep_NJ.js → index.esm-Dy6Z9Ung.js} +1 -1
- zenml/zen_server/dashboard/assets/{kubernetes--g7r02Zu.js → kubernetes-B2wmAJ1d.js} +1 -1
- zenml/zen_server/dashboard/assets/{layout-DCSYN7-C.js → layout-BtHBmE4w.js} +1 -1
- zenml/zen_server/dashboard/assets/{link-external-CBEk6kEG.js → link-external-b9AXw_sW.js} +1 -1
- zenml/zen_server/dashboard/assets/{login-mutation-DTcAFP1l.js → login-mutation-hf-lK87O.js} +1 -1
- zenml/zen_server/dashboard/assets/{logs-D5bdJGur.js → logs-WMSM52RF.js} +1 -1
- zenml/zen_server/dashboard/assets/{not-found-Cc-JkRH2.js → not-found-BGirLjU-.js} +1 -1
- zenml/zen_server/dashboard/assets/{package-Cs35Szwh.js → package-C6uypY4h.js} +1 -1
- zenml/zen_server/dashboard/assets/page-0JE_-Ec1.js +1 -0
- zenml/zen_server/dashboard/assets/{page-DH_Z7iW1.js → page-6m6yHHlE.js} +1 -1
- zenml/zen_server/dashboard/assets/page-BDigxVpo.js +1 -0
- zenml/zen_server/dashboard/assets/page-BR68V0V1.js +1 -0
- zenml/zen_server/dashboard/assets/page-BRLpxOt0.js +1 -0
- zenml/zen_server/dashboard/assets/{page-BQQKaabe.js → page-BU7huvKw.js} +3 -3
- zenml/zen_server/dashboard/assets/page-BvqLv2Ky.js +1 -0
- zenml/zen_server/dashboard/assets/page-C00YAkaB.js +1 -0
- zenml/zen_server/dashboard/assets/{page-N4qoPHKb.js → page-CD-DcWoy.js} +1 -1
- zenml/zen_server/dashboard/assets/page-COXXJj1k.js +1 -0
- zenml/zen_server/dashboard/assets/page-CbpvrsDL.js +1 -0
- zenml/zen_server/dashboard/assets/page-CdMWnQak.js +1 -0
- zenml/zen_server/dashboard/assets/{page-ClUVkl-O.js → page-CjGdWY13.js} +1 -1
- zenml/zen_server/dashboard/assets/page-CwxrFarU.js +1 -0
- zenml/zen_server/dashboard/assets/{page-DLixvR-7.js → page-D01JhjQB.js} +1 -1
- zenml/zen_server/dashboard/assets/page-D6uU2ax4.js +1 -0
- zenml/zen_server/dashboard/assets/page-D7S3aCbF.js +1 -0
- zenml/zen_server/dashboard/assets/{page-9yplj5JT.js → page-DLC-bNBP.js} +1 -1
- zenml/zen_server/dashboard/assets/page-DXSTpqRD.js +1 -0
- zenml/zen_server/dashboard/assets/{page-DzpVUZ8f.js → page-DakHVWXF.js} +1 -1
- zenml/zen_server/dashboard/assets/{page-DIOXwhiD.js → page-Df-Fw0aq.js} +1 -1
- zenml/zen_server/dashboard/assets/{page-B-y2XKIc.js → page-DfbXf_8s.js} +1 -1
- zenml/zen_server/dashboard/assets/page-DjRJCGb3.js +1 -0
- zenml/zen_server/dashboard/assets/{page-C0N5q3l7.js → page-Djikxq_S.js} +1 -1
- zenml/zen_server/dashboard/assets/page-Dnovpa0i.js +3 -0
- zenml/zen_server/dashboard/assets/page-Dot3LPmL.js +1 -0
- zenml/zen_server/dashboard/assets/page-Vcxara9U.js +1 -0
- zenml/zen_server/dashboard/assets/page-Xynx4btY.js +14 -0
- zenml/zen_server/dashboard/assets/page-YpKAqVSa.js +1 -0
- zenml/zen_server/dashboard/assets/page-yYC9OI-E.js +1 -0
- zenml/zen_server/dashboard/assets/{persist-DNb5cdrU.js → persist-Coz7ZWvz.js} +1 -1
- zenml/zen_server/dashboard/assets/{persist-CP0JmYZ4.js → persist-GjC8PZoC.js} +1 -1
- zenml/zen_server/dashboard/assets/{plus-C9IxgN2M.js → plus-tf1V2hTJ.js} +1 -1
- zenml/zen_server/dashboard/assets/{refresh-BVu22P_C.js → refresh-BjOeWlEq.js} +1 -1
- zenml/zen_server/dashboard/assets/{rocket-CONEmRmB.js → rocket-DjT2cDvG.js} +1 -1
- zenml/zen_server/dashboard/assets/sharedSchema-CQb14VSr.js +14 -0
- zenml/zen_server/dashboard/assets/stack-detail-query-OPEW-cDJ.js +1 -0
- zenml/zen_server/dashboard/assets/{tick-circle-CM1ZScbQ.js → tick-circle-BEX_Tp4v.js} +1 -1
- zenml/zen_server/dashboard/assets/{trash-DkJHMOg7.js → trash-arLUMWMS.js} +1 -1
- zenml/zen_server/dashboard/assets/{update-server-settings-mutation-DsU8cNVl.js → update-server-settings-mutation-LwuQfHYn.js} +1 -1
- zenml/zen_server/dashboard/assets/upgrade-form-CwRHBuXB.webp +0 -0
- zenml/zen_server/dashboard/assets/url-CkvKAnwF.js +1 -0
- zenml/zen_server/dashboard/assets/{zod-D89GC_vc.js → zod-BwEbpOxH.js} +1 -1
- zenml/zen_server/dashboard/index.html +7 -7
- zenml/zen_server/deploy/helm/Chart.yaml +1 -1
- zenml/zen_server/deploy/helm/README.md +2 -2
- zenml/zen_server/exceptions.py +2 -0
- zenml/zen_server/jwt.py +30 -13
- zenml/zen_server/rbac/endpoint_utils.py +43 -1
- zenml/zen_server/rbac/utils.py +0 -2
- zenml/zen_server/routers/artifact_version_endpoints.py +27 -1
- zenml/zen_server/routers/auth_endpoints.py +134 -102
- zenml/zen_server/routers/logs_endpoints.py +66 -0
- zenml/zen_server/routers/workspaces_endpoints.py +3 -4
- zenml/zen_server/template_execution/utils.py +14 -16
- zenml/zen_server/utils.py +27 -0
- zenml/zen_server/zen_server_api.py +6 -3
- zenml/zen_stores/migrations/versions/0.70.0_release.py +23 -0
- zenml/zen_stores/migrations/versions/1cb6477f72d6_move_artifact_save_type.py +99 -0
- zenml/zen_stores/migrations/versions/904464ea4041_add_pipeline_model_run_unique_constraints.py +192 -0
- zenml/zen_stores/migrations/versions/b557b2871693_update_step_run_input_types.py +33 -0
- zenml/zen_stores/rest_zen_store.py +68 -64
- zenml/zen_stores/schemas/artifact_schemas.py +8 -1
- zenml/zen_stores/schemas/model_schemas.py +27 -3
- zenml/zen_stores/schemas/pipeline_run_schemas.py +6 -1
- zenml/zen_stores/schemas/pipeline_schemas.py +8 -2
- zenml/zen_stores/schemas/run_metadata_schemas.py +1 -48
- zenml/zen_stores/schemas/step_run_schemas.py +18 -10
- zenml/zen_stores/sql_zen_store.py +283 -219
- zenml/zen_stores/zen_store_interface.py +15 -42
- {zenml_nightly-0.68.1.dev20241103.dist-info → zenml_nightly-0.70.0.dev20241115.dist-info}/METADATA +2 -2
- {zenml_nightly-0.68.1.dev20241103.dist-info → zenml_nightly-0.70.0.dev20241115.dist-info}/RECORD +226 -203
- zenml/zen_server/dashboard/assets/AlertDialogDropdownItem-C6N2rGrB.js +0 -1
- zenml/zen_server/dashboard/assets/ComponentBadge-DUiEYJHu.js +0 -1
- zenml/zen_server/dashboard/assets/ComponentFallbackDialog-BFoH5K4V.js +0 -1
- zenml/zen_server/dashboard/assets/ComponentIcon-CAIoUis2.js +0 -1
- zenml/zen_server/dashboard/assets/Partials-YPBB3V4q.js +0 -1
- zenml/zen_server/dashboard/assets/ProviderIcon-Bb3Xha5A.js +0 -1
- zenml/zen_server/dashboard/assets/RunSelector-DCiL3M1c.js +0 -1
- zenml/zen_server/dashboard/assets/SearchField-DfUiGFVd.js +0 -1
- zenml/zen_server/dashboard/assets/Tick-CykQFPj2.js +0 -1
- zenml/zen_server/dashboard/assets/cloud-only-B-s_HMDm.js +0 -1
- zenml/zen_server/dashboard/assets/codespaces-BitYDX9d.gif +0 -0
- zenml/zen_server/dashboard/assets/create-stack-CEmaPZ4c.js +0 -1
- zenml/zen_server/dashboard/assets/delete-run-D-LKbGyz.js +0 -1
- zenml/zen_server/dashboard/assets/index-Bpmj40BI.js +0 -1
- zenml/zen_server/dashboard/assets/index-CbU4Ln_E.css +0 -1
- zenml/zen_server/dashboard/assets/index-DKPhqP2B.js +0 -1
- zenml/zen_server/dashboard/assets/page-BBpOxVcY.js +0 -1
- zenml/zen_server/dashboard/assets/page-BRInM1Lg.js +0 -1
- zenml/zen_server/dashboard/assets/page-BjjlMk7s.js +0 -1
- zenml/zen_server/dashboard/assets/page-Bvd7YH2A.js +0 -1
- zenml/zen_server/dashboard/assets/page-CT3Nep8W.js +0 -1
- zenml/zen_server/dashboard/assets/page-C_f47pBf.js +0 -1
- zenml/zen_server/dashboard/assets/page-Cmv8C_yM.js +0 -3
- zenml/zen_server/dashboard/assets/page-CyN2bdWG.js +0 -1
- zenml/zen_server/dashboard/assets/page-CzzXH4fs.js +0 -1
- zenml/zen_server/dashboard/assets/page-DTlGjgnG.js +0 -1
- zenml/zen_server/dashboard/assets/page-Dbpl86h0.js +0 -1
- zenml/zen_server/dashboard/assets/page-Ddgy6kDS.js +0 -1
- zenml/zen_server/dashboard/assets/page-DtCAfBLy.js +0 -9
- zenml/zen_server/dashboard/assets/page-Dx16z7nA.js +0 -1
- zenml/zen_server/dashboard/assets/page-McUyYbo1.js +0 -1
- zenml/zen_server/dashboard/assets/page-T1P3RyAR.js +0 -1
- zenml/zen_server/dashboard/assets/page-bKaULTGG.js +0 -1
- zenml/zen_server/dashboard/assets/page-sbXUJy9t.js +0 -1
- zenml/zen_server/dashboard/assets/sharedSchema-TMLu-nYQ.js +0 -14
- zenml/zen_server/dashboard/assets/stack-detail-query-xmYxSsUY.js +0 -1
- zenml/zen_server/dashboard/assets/url-D5le3J4q.js +0 -1
- zenml/zen_server/routers/run_metadata_endpoints.py +0 -96
- {zenml_nightly-0.68.1.dev20241103.dist-info → zenml_nightly-0.70.0.dev20241115.dist-info}/LICENSE +0 -0
- {zenml_nightly-0.68.1.dev20241103.dist-info → zenml_nightly-0.70.0.dev20241115.dist-info}/WHEEL +0 -0
- {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
|
-
|
50
|
-
|
51
|
-
|
52
|
-
|
53
|
-
|
54
|
-
|
55
|
-
|
56
|
-
|
57
|
-
|
58
|
-
|
59
|
-
|
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
|
-
|
68
|
-
|
69
|
-
|
70
|
-
|
71
|
-
|
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
|
-
|
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
|
-
|
50
|
-
|
51
|
-
|
52
|
-
|
53
|
-
|
54
|
-
|
55
|
-
|
56
|
-
|
57
|
-
|
58
|
-
|
59
|
-
|
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
|
-
|
68
|
-
|
69
|
-
|
70
|
-
|
71
|
-
|
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
|
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
|
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,
|
@@ -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
|
27
|
+
from langchain_community.embeddings import (
|
28
|
+
OpenAIEmbeddings,
|
29
|
+
)
|
28
30
|
|
29
31
|
|
30
32
|
class LangchainOpenaiEmbeddingMaterializer(CloudpickleMaterializer):
|
31
|
-
"""
|
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
|
-
|
47
|
-
|
48
|
-
|
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
|
-
|
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
|
-
|
50
|
-
|
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
|
-
|
54
|
-
|
55
|
-
|
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
|
-
|
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
|
-
|
69
|
-
|
70
|
-
|
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
|
-
|
74
|
-
|
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
|
-
|
395
|
-
|
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
|
-
|
491
|
-
|
492
|
-
|
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(
|
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
|
-
|
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")
|
@@ -15,8 +15,9 @@
|
|
15
15
|
|
16
16
|
import io
|
17
17
|
import sys
|
18
|
+
from typing import Optional
|
18
19
|
|
19
|
-
import
|
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(
|
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
|
-
|
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)}')
|
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
|
-
|
75
|
-
|
76
|
-
|
77
|
-
|
78
|
-
message
|
79
|
-
|
80
|
-
|
81
|
-
|
82
|
-
|
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-
|
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
|
-
|
61
|
-
|
62
|
-
|
63
|
-
|
64
|
-
|
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
|
-
|
68
|
-
|
69
|
-
|
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
|
-
|
78
|
-
|
79
|
-
|
80
|
-
|
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
|
-
|
83
|
-
|
80
|
+
# save the image in a temporary directory
|
81
|
+
image.save(temp_image_path)
|
84
82
|
|
85
|
-
|
86
|
-
|
87
|
-
|
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
|
-
|
49
|
-
|
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
|
-
|
52
|
-
|
53
|
-
|
54
|
-
|
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
|
-
|
71
|
-
|
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
|
-
|
74
|
-
|
65
|
+
# Convert the table to a Polars data frame or series
|
66
|
+
data = pl.from_arrow(table)
|
75
67
|
|
76
|
-
|
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
|
-
|
111
|
-
|
112
|
-
|
113
|
-
|
114
|
-
|
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)
|