zenml-nightly 0.68.1.dev20241106__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/utils.py +138 -79
- 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/aws/orchestrators/sagemaker_orchestrator.py +20 -18
- zenml/integrations/azure/orchestrators/azureml_orchestrator.py +1 -1
- 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/lightning/orchestrators/lightning_orchestrator.py +29 -9
- zenml/integrations/tensorboard/visualizers/tensorboard_visualizer.py +60 -54
- zenml/integrations/vllm/services/vllm_deployment.py +16 -7
- 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 +21 -13
- zenml/orchestrators/utils.py +45 -26
- 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 +3 -2
- zenml/steps/utils.py +13 -2
- 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.dev20241106.dist-info → zenml_nightly-0.70.0.dev20241115.dist-info}/METADATA +2 -2
- {zenml_nightly-0.68.1.dev20241106.dist-info → zenml_nightly-0.70.0.dev20241115.dist-info}/RECORD +194 -179
- 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.dev20241106.dist-info → zenml_nightly-0.70.0.dev20241115.dist-info}/LICENSE +0 -0
- {zenml_nightly-0.68.1.dev20241106.dist-info → zenml_nightly-0.70.0.dev20241115.dist-info}/WHEEL +0 -0
- {zenml_nightly-0.68.1.dev20241106.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
|
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
|
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
|
-
|
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
|
-
|
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
|
-
|
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
|
5317
|
-
# Creating the run failed
|
5318
|
-
# -
|
5319
|
-
#
|
5320
|
-
#
|
5321
|
-
#
|
5322
|
-
#
|
5323
|
-
#
|
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
|
5341
|
-
|
5342
|
-
|
5343
|
-
|
5344
|
-
|
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
|
-
|
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=
|
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,
|
8188
|
-
|
8189
|
-
|
8190
|
-
|
8191
|
-
|
8192
|
-
|
8193
|
-
|
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
|
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
|
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
|
-
|
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
|
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
|
-
|
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
|
-
|
10232
|
-
|
10233
|
-
|
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
|
-
|
10239
|
-
all_versions.number + 1 if all_versions else 1
|
10240
|
-
)
|
10259
|
+
model_version_id = None
|
10241
10260
|
|
10242
|
-
|
10243
|
-
|
10244
|
-
|
10245
|
-
|
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
|
-
|
10248
|
-
|
10249
|
-
|
10250
|
-
|
10275
|
+
model_version_schema = ModelVersionSchema.from_request(
|
10276
|
+
model_version
|
10277
|
+
)
|
10278
|
+
session.add(model_version_schema)
|
10279
|
+
session.commit()
|
10251
10280
|
|
10252
|
-
|
10253
|
-
|
10254
|
-
|
10255
|
-
|
10256
|
-
|
10257
|
-
)
|
10258
|
-
|
10259
|
-
|
10260
|
-
|
10261
|
-
|
10262
|
-
|
10263
|
-
|
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
|
-
|
10266
|
-
|
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:
|