zenml-nightly 0.68.1.dev20241112__py3-none-any.whl → 0.70.0.dev20241114__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 (186) 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/utils.py +5 -1
  6. zenml/client.py +6 -1
  7. zenml/config/server_config.py +17 -1
  8. zenml/constants.py +1 -7
  9. zenml/enums.py +8 -0
  10. zenml/model/utils.py +5 -0
  11. zenml/models/v2/base/filter.py +121 -8
  12. zenml/models/v2/core/artifact_version.py +23 -0
  13. zenml/models/v2/core/model_version.py +23 -0
  14. zenml/models/v2/core/pipeline_run.py +22 -1
  15. zenml/models/v2/core/step_run.py +22 -0
  16. zenml/orchestrators/base_orchestrator.py +12 -1
  17. zenml/orchestrators/step_launcher.py +2 -1
  18. zenml/orchestrators/utils.py +45 -26
  19. zenml/steps/utils.py +5 -0
  20. zenml/utils/metadata_utils.py +335 -0
  21. zenml/zen_server/auth.py +221 -3
  22. zenml/zen_server/cache.py +208 -0
  23. zenml/zen_server/dashboard/assets/{404-DT4QRUqN.js → 404-NVXKFp-x.js} +1 -1
  24. zenml/zen_server/dashboard/assets/{@radix-DP6vWzyx.js → @radix-DeK6qiuw.js} +1 -1
  25. zenml/zen_server/dashboard/assets/{@react-router-BMhZulnd.js → @react-router-B3Z5rLr2.js} +1 -1
  26. zenml/zen_server/dashboard/assets/{@reactflow-8U9qNlMR.js → @reactflow-CK0KJUen.js} +2 -2
  27. zenml/zen_server/dashboard/assets/{@tanstack-BUCbhJyH.js → @tanstack-DT5WLu9C.js} +1 -1
  28. zenml/zen_server/dashboard/assets/AlertDialogDropdownItem-DezXKmDf.js +1 -0
  29. zenml/zen_server/dashboard/assets/{CodeSnippet-CqybNv0k.js → CodeSnippet-JzR8CEtw.js} +2 -2
  30. zenml/zen_server/dashboard/assets/{CollapsibleCard-0r_8G2Lj.js → CollapsibleCard-DQW_ktMO.js} +1 -1
  31. zenml/zen_server/dashboard/assets/{Commands-BDjgBQKi.js → Commands-DL2kwkRd.js} +1 -1
  32. zenml/zen_server/dashboard/assets/ComponentBadge-D_g62Wv8.js +1 -0
  33. zenml/zen_server/dashboard/assets/{CopyButton-C745BrKi.js → CopyButton-LNcWaa14.js} +1 -1
  34. zenml/zen_server/dashboard/assets/{CsvVizualization-PpAq0CeZ.js → CsvVizualization-DknpE5ej.js} +5 -5
  35. zenml/zen_server/dashboard/assets/{DialogItem-DcVCZEew.js → DialogItem-Bxf8FuAT.js} +1 -1
  36. zenml/zen_server/dashboard/assets/{DisplayDate-BeXgUG_C.js → DisplayDate-CDMUcQHS.js} +1 -1
  37. zenml/zen_server/dashboard/assets/{EmptyState-DeK7H4pr.js → EmptyState-BzdlCwp3.js} +1 -1
  38. zenml/zen_server/dashboard/assets/{Error-BMlzibXj.js → Error-DYflYyps.js} +1 -1
  39. zenml/zen_server/dashboard/assets/ExecutionStatus-C7zyIQKZ.js +1 -0
  40. zenml/zen_server/dashboard/assets/{Helpbox-BLf40fLV.js → Helpbox-oYSGpLqd.js} +1 -1
  41. zenml/zen_server/dashboard/assets/{Infobox-BwisKifi.js → Infobox-Cx4xGoXR.js} +1 -1
  42. zenml/zen_server/dashboard/assets/{InlineAvatar-jEgodSgX.js → InlineAvatar-DiGOWNKF.js} +1 -1
  43. zenml/zen_server/dashboard/assets/{Lock-3lLt1ih0.js → Lock-CYYy18Mm.js} +1 -1
  44. zenml/zen_server/dashboard/assets/{MarkdownVisualization-8O9kTr-2.js → MarkdownVisualization-ylXaAxev.js} +1 -1
  45. zenml/zen_server/dashboard/assets/NestedCollapsible-DYbgyKxK.js +1 -0
  46. zenml/zen_server/dashboard/assets/{NumberBox-T9eELfLZ.js → NumberBox-Dtp3J6g5.js} +1 -1
  47. zenml/zen_server/dashboard/assets/Partials-03iZf8-N.js +1 -0
  48. zenml/zen_server/dashboard/assets/{PasswordChecker-CW0kqY0W.js → PasswordChecker-B0nadgh6.js} +1 -1
  49. zenml/zen_server/dashboard/assets/ProBadge-D_EB8HNo.js +1 -0
  50. zenml/zen_server/dashboard/assets/ProCta-DqNS4v3x.js +1 -0
  51. zenml/zen_server/dashboard/assets/ProviderIcon-Bki2aw8w.js +1 -0
  52. zenml/zen_server/dashboard/assets/{ProviderRadio-BROY1700.js → ProviderRadio-8f43sPD4.js} +1 -1
  53. zenml/zen_server/dashboard/assets/RunSelector-DkPiIiNr.js +1 -0
  54. zenml/zen_server/dashboard/assets/RunsBody-07YEO7qI.js +1 -0
  55. zenml/zen_server/dashboard/assets/SearchField-lp1KgU4e.js +1 -0
  56. zenml/zen_server/dashboard/assets/{SecretTooltip-C_qByGWB.js → SecretTooltip-CgnbyeOx.js} +1 -1
  57. zenml/zen_server/dashboard/assets/{SetPassword-7pRB00El.js → SetPassword-CpP418A2.js} +1 -1
  58. zenml/zen_server/dashboard/assets/StackList-WvuKQusZ.js +1 -0
  59. zenml/zen_server/dashboard/assets/Tabs-BktHkCJJ.js +1 -0
  60. zenml/zen_server/dashboard/assets/Tick-BlMoIlJT.js +1 -0
  61. zenml/zen_server/dashboard/assets/{UpdatePasswordSchemas-DckMEkFf.js → UpdatePasswordSchemas-Sc0A0pP-.js} +1 -1
  62. zenml/zen_server/dashboard/assets/{UsageReason-DVceN14P.js → UsageReason-YYduL4fj.js} +1 -1
  63. zenml/zen_server/dashboard/assets/{WizardFooter-CW0Cvd70.js → WizardFooter-dgmizSJC.js} +1 -1
  64. zenml/zen_server/dashboard/assets/all-pipeline-runs-query-D-c2G6lV.js +1 -0
  65. zenml/zen_server/dashboard/assets/check-DloQpStc.js +1 -0
  66. zenml/zen_server/dashboard/assets/{check-circle-Dwxliy1Z.js → check-circle-jNbX5-sR.js} +1 -1
  67. zenml/zen_server/dashboard/assets/{chevron-down-8wLBS5pQ.js → chevron-down-6JyMkfjR.js} +1 -1
  68. zenml/zen_server/dashboard/assets/{chevron-right-double-DoD8iXWM.js → chevron-right-double-D7ojK9Co.js} +1 -1
  69. zenml/zen_server/dashboard/assets/{code-browser-CZUQs3Wa.js → code-browser-CUFUIHfp.js} +1 -1
  70. zenml/zen_server/dashboard/assets/{copy-CaSMXwiU.js → copy-C8XQA2Ug.js} +1 -1
  71. zenml/zen_server/dashboard/assets/create-stack-DM_JPgef.js +1 -0
  72. zenml/zen_server/dashboard/assets/delete-run-CJdh1P_h.js +1 -0
  73. zenml/zen_server/dashboard/assets/{docker-BFNgg-z3.js → docker-BdA9vrnW.js} +1 -1
  74. zenml/zen_server/dashboard/assets/{dots-horizontal-DK5Duzx4.js → dots-horizontal-otGBOSDJ.js} +1 -1
  75. zenml/zen_server/dashboard/assets/{form-schemas-1AyOCx90.js → form-schemas-K6FYKjwa.js} +1 -1
  76. zenml/zen_server/dashboard/assets/{gcp-7M2Yf3ZK.js → gcp-CFtm4BA7.js} +1 -1
  77. zenml/zen_server/dashboard/assets/{help-Dam461dC.js → help-Cc9bBIJH.js} +1 -1
  78. zenml/zen_server/dashboard/assets/index-B1mVPYxf.js +1 -0
  79. zenml/zen_server/dashboard/assets/index-BAkC7FXi.js +1 -0
  80. zenml/zen_server/dashboard/assets/{index-QQb7wQEC.js → index-CCOPpudF.js} +8 -8
  81. zenml/zen_server/dashboard/assets/index-CEV4Cvaf.js +1 -0
  82. zenml/zen_server/dashboard/assets/index-DlGvJQPn.css +1 -0
  83. zenml/zen_server/dashboard/assets/{index-BVJ8n2-j.js → index-Uu49AX48.js} +1 -1
  84. zenml/zen_server/dashboard/assets/{index.esm-cuVep_NJ.js → index.esm-Dy6Z9Ung.js} +1 -1
  85. zenml/zen_server/dashboard/assets/{kubernetes--g7r02Zu.js → kubernetes-B2wmAJ1d.js} +1 -1
  86. zenml/zen_server/dashboard/assets/{layout-DCSYN7-C.js → layout-BtHBmE4w.js} +1 -1
  87. zenml/zen_server/dashboard/assets/{link-external-CBEk6kEG.js → link-external-b9AXw_sW.js} +1 -1
  88. zenml/zen_server/dashboard/assets/{login-mutation-DTcAFP1l.js → login-mutation-hf-lK87O.js} +1 -1
  89. zenml/zen_server/dashboard/assets/{logs-D5bdJGur.js → logs-WMSM52RF.js} +1 -1
  90. zenml/zen_server/dashboard/assets/{not-found-Cc-JkRH2.js → not-found-BGirLjU-.js} +1 -1
  91. zenml/zen_server/dashboard/assets/{package-Cs35Szwh.js → package-C6uypY4h.js} +1 -1
  92. zenml/zen_server/dashboard/assets/page-0JE_-Ec1.js +1 -0
  93. zenml/zen_server/dashboard/assets/{page-DH_Z7iW1.js → page-6m6yHHlE.js} +1 -1
  94. zenml/zen_server/dashboard/assets/page-BDigxVpo.js +1 -0
  95. zenml/zen_server/dashboard/assets/page-BR68V0V1.js +1 -0
  96. zenml/zen_server/dashboard/assets/page-BRLpxOt0.js +1 -0
  97. zenml/zen_server/dashboard/assets/{page-BQQKaabe.js → page-BU7huvKw.js} +3 -3
  98. zenml/zen_server/dashboard/assets/page-BvqLv2Ky.js +1 -0
  99. zenml/zen_server/dashboard/assets/page-C00YAkaB.js +1 -0
  100. zenml/zen_server/dashboard/assets/{page-N4qoPHKb.js → page-CD-DcWoy.js} +1 -1
  101. zenml/zen_server/dashboard/assets/page-COXXJj1k.js +1 -0
  102. zenml/zen_server/dashboard/assets/page-CbpvrsDL.js +1 -0
  103. zenml/zen_server/dashboard/assets/page-CdMWnQak.js +1 -0
  104. zenml/zen_server/dashboard/assets/{page-ClUVkl-O.js → page-CjGdWY13.js} +1 -1
  105. zenml/zen_server/dashboard/assets/page-CwxrFarU.js +1 -0
  106. zenml/zen_server/dashboard/assets/{page-DLixvR-7.js → page-D01JhjQB.js} +1 -1
  107. zenml/zen_server/dashboard/assets/page-D6uU2ax4.js +1 -0
  108. zenml/zen_server/dashboard/assets/page-D7S3aCbF.js +1 -0
  109. zenml/zen_server/dashboard/assets/{page-9yplj5JT.js → page-DLC-bNBP.js} +1 -1
  110. zenml/zen_server/dashboard/assets/page-DXSTpqRD.js +1 -0
  111. zenml/zen_server/dashboard/assets/{page-DzpVUZ8f.js → page-DakHVWXF.js} +1 -1
  112. zenml/zen_server/dashboard/assets/{page-DIOXwhiD.js → page-Df-Fw0aq.js} +1 -1
  113. zenml/zen_server/dashboard/assets/{page-B-y2XKIc.js → page-DfbXf_8s.js} +1 -1
  114. zenml/zen_server/dashboard/assets/page-DjRJCGb3.js +1 -0
  115. zenml/zen_server/dashboard/assets/{page-C0N5q3l7.js → page-Djikxq_S.js} +1 -1
  116. zenml/zen_server/dashboard/assets/page-Dnovpa0i.js +3 -0
  117. zenml/zen_server/dashboard/assets/page-Dot3LPmL.js +1 -0
  118. zenml/zen_server/dashboard/assets/page-Vcxara9U.js +1 -0
  119. zenml/zen_server/dashboard/assets/page-Xynx4btY.js +14 -0
  120. zenml/zen_server/dashboard/assets/page-YpKAqVSa.js +1 -0
  121. zenml/zen_server/dashboard/assets/page-yYC9OI-E.js +1 -0
  122. zenml/zen_server/dashboard/assets/{persist-DNb5cdrU.js → persist-Coz7ZWvz.js} +1 -1
  123. zenml/zen_server/dashboard/assets/{persist-CP0JmYZ4.js → persist-GjC8PZoC.js} +1 -1
  124. zenml/zen_server/dashboard/assets/{plus-C9IxgN2M.js → plus-tf1V2hTJ.js} +1 -1
  125. zenml/zen_server/dashboard/assets/{refresh-BVu22P_C.js → refresh-BjOeWlEq.js} +1 -1
  126. zenml/zen_server/dashboard/assets/{rocket-CONEmRmB.js → rocket-DjT2cDvG.js} +1 -1
  127. zenml/zen_server/dashboard/assets/sharedSchema-CQb14VSr.js +14 -0
  128. zenml/zen_server/dashboard/assets/stack-detail-query-OPEW-cDJ.js +1 -0
  129. zenml/zen_server/dashboard/assets/{tick-circle-CM1ZScbQ.js → tick-circle-BEX_Tp4v.js} +1 -1
  130. zenml/zen_server/dashboard/assets/{trash-DkJHMOg7.js → trash-arLUMWMS.js} +1 -1
  131. zenml/zen_server/dashboard/assets/{update-server-settings-mutation-DsU8cNVl.js → update-server-settings-mutation-LwuQfHYn.js} +1 -1
  132. zenml/zen_server/dashboard/assets/upgrade-form-CwRHBuXB.webp +0 -0
  133. zenml/zen_server/dashboard/assets/url-CkvKAnwF.js +1 -0
  134. zenml/zen_server/dashboard/assets/{zod-D89GC_vc.js → zod-BwEbpOxH.js} +1 -1
  135. zenml/zen_server/dashboard/index.html +7 -7
  136. zenml/zen_server/deploy/helm/Chart.yaml +1 -1
  137. zenml/zen_server/deploy/helm/README.md +2 -2
  138. zenml/zen_server/jwt.py +30 -13
  139. zenml/zen_server/routers/auth_endpoints.py +134 -102
  140. zenml/zen_server/template_execution/utils.py +14 -16
  141. zenml/zen_server/utils.py +27 -0
  142. zenml/zen_server/zen_server_api.py +4 -1
  143. zenml/zen_stores/migrations/versions/0.70.0_release.py +23 -0
  144. zenml/zen_stores/rest_zen_store.py +13 -10
  145. {zenml_nightly-0.68.1.dev20241112.dist-info → zenml_nightly-0.70.0.dev20241114.dist-info}/METADATA +2 -2
  146. {zenml_nightly-0.68.1.dev20241112.dist-info → zenml_nightly-0.70.0.dev20241114.dist-info}/RECORD +149 -137
  147. zenml/zen_server/dashboard/assets/AlertDialogDropdownItem-C6N2rGrB.js +0 -1
  148. zenml/zen_server/dashboard/assets/ComponentBadge-DUiEYJHu.js +0 -1
  149. zenml/zen_server/dashboard/assets/ComponentFallbackDialog-BFoH5K4V.js +0 -1
  150. zenml/zen_server/dashboard/assets/ComponentIcon-CAIoUis2.js +0 -1
  151. zenml/zen_server/dashboard/assets/Partials-YPBB3V4q.js +0 -1
  152. zenml/zen_server/dashboard/assets/ProviderIcon-Bb3Xha5A.js +0 -1
  153. zenml/zen_server/dashboard/assets/RunSelector-DCiL3M1c.js +0 -1
  154. zenml/zen_server/dashboard/assets/SearchField-DfUiGFVd.js +0 -1
  155. zenml/zen_server/dashboard/assets/Tick-CykQFPj2.js +0 -1
  156. zenml/zen_server/dashboard/assets/cloud-only-B-s_HMDm.js +0 -1
  157. zenml/zen_server/dashboard/assets/codespaces-BitYDX9d.gif +0 -0
  158. zenml/zen_server/dashboard/assets/create-stack-CEmaPZ4c.js +0 -1
  159. zenml/zen_server/dashboard/assets/delete-run-D-LKbGyz.js +0 -1
  160. zenml/zen_server/dashboard/assets/index-Bpmj40BI.js +0 -1
  161. zenml/zen_server/dashboard/assets/index-CbU4Ln_E.css +0 -1
  162. zenml/zen_server/dashboard/assets/index-DKPhqP2B.js +0 -1
  163. zenml/zen_server/dashboard/assets/page-BBpOxVcY.js +0 -1
  164. zenml/zen_server/dashboard/assets/page-BRInM1Lg.js +0 -1
  165. zenml/zen_server/dashboard/assets/page-BjjlMk7s.js +0 -1
  166. zenml/zen_server/dashboard/assets/page-Bvd7YH2A.js +0 -1
  167. zenml/zen_server/dashboard/assets/page-CT3Nep8W.js +0 -1
  168. zenml/zen_server/dashboard/assets/page-C_f47pBf.js +0 -1
  169. zenml/zen_server/dashboard/assets/page-Cmv8C_yM.js +0 -3
  170. zenml/zen_server/dashboard/assets/page-CyN2bdWG.js +0 -1
  171. zenml/zen_server/dashboard/assets/page-CzzXH4fs.js +0 -1
  172. zenml/zen_server/dashboard/assets/page-DTlGjgnG.js +0 -1
  173. zenml/zen_server/dashboard/assets/page-Dbpl86h0.js +0 -1
  174. zenml/zen_server/dashboard/assets/page-Ddgy6kDS.js +0 -1
  175. zenml/zen_server/dashboard/assets/page-DtCAfBLy.js +0 -9
  176. zenml/zen_server/dashboard/assets/page-Dx16z7nA.js +0 -1
  177. zenml/zen_server/dashboard/assets/page-McUyYbo1.js +0 -1
  178. zenml/zen_server/dashboard/assets/page-T1P3RyAR.js +0 -1
  179. zenml/zen_server/dashboard/assets/page-bKaULTGG.js +0 -1
  180. zenml/zen_server/dashboard/assets/page-sbXUJy9t.js +0 -1
  181. zenml/zen_server/dashboard/assets/sharedSchema-TMLu-nYQ.js +0 -14
  182. zenml/zen_server/dashboard/assets/stack-detail-query-xmYxSsUY.js +0 -1
  183. zenml/zen_server/dashboard/assets/url-D5le3J4q.js +0 -1
  184. {zenml_nightly-0.68.1.dev20241112.dist-info → zenml_nightly-0.70.0.dev20241114.dist-info}/LICENSE +0 -0
  185. {zenml_nightly-0.68.1.dev20241112.dist-info → zenml_nightly-0.70.0.dev20241114.dist-info}/WHEEL +0 -0
  186. {zenml_nightly-0.68.1.dev20241112.dist-info → zenml_nightly-0.70.0.dev20241114.dist-info}/entry_points.txt +0 -0
@@ -587,6 +587,7 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
587
587
  "stack_component",
588
588
  "pipeline_name",
589
589
  "templatable",
590
+ "run_metadata",
590
591
  ]
591
592
  name: Optional[str] = Field(
592
593
  default=None,
@@ -665,6 +666,10 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
665
666
  default=None,
666
667
  description="Name/ID of the user that created the run.",
667
668
  )
669
+ run_metadata: Optional[Dict[str, str]] = Field(
670
+ default=None,
671
+ description="The run_metadata to filter the pipeline runs by.",
672
+ )
668
673
  # TODO: Remove once frontend is ready for it. This is replaced by the more
669
674
  # generic `pipeline` filter below.
670
675
  pipeline_name: Optional[str] = Field(
@@ -694,7 +699,6 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
694
699
  templatable: Optional[bool] = Field(
695
700
  default=None, description="Whether the run is templatable."
696
701
  )
697
-
698
702
  model_config = ConfigDict(protected_namespaces=())
699
703
 
700
704
  def get_custom_filters(
@@ -718,6 +722,7 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
718
722
  PipelineDeploymentSchema,
719
723
  PipelineRunSchema,
720
724
  PipelineSchema,
725
+ RunMetadataSchema,
721
726
  ScheduleSchema,
722
727
  StackComponentSchema,
723
728
  StackCompositionSchema,
@@ -887,5 +892,21 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
887
892
  )
888
893
 
889
894
  custom_filters.append(templatable_filter)
895
+ if self.run_metadata is not None:
896
+ from zenml.enums import MetadataResourceTypes
897
+
898
+ for key, value in self.run_metadata.items():
899
+ additional_filter = and_(
900
+ RunMetadataSchema.resource_id == PipelineRunSchema.id,
901
+ RunMetadataSchema.resource_type
902
+ == MetadataResourceTypes.PIPELINE_RUN,
903
+ RunMetadataSchema.key == key,
904
+ self.generate_custom_query_conditions_for_column(
905
+ value=value,
906
+ table=RunMetadataSchema,
907
+ column="value",
908
+ ),
909
+ )
910
+ custom_filters.append(additional_filter)
890
911
 
891
912
  return custom_filters
@@ -509,6 +509,7 @@ class StepRunFilter(WorkspaceScopedFilter):
509
509
  FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
510
510
  *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
511
511
  "model",
512
+ "run_metadata",
512
513
  ]
513
514
 
514
515
  name: Optional[str] = Field(
@@ -571,6 +572,10 @@ class StepRunFilter(WorkspaceScopedFilter):
571
572
  default=None,
572
573
  description="Name/ID of the model associated with the step run.",
573
574
  )
575
+ run_metadata: Optional[Dict[str, str]] = Field(
576
+ default=None,
577
+ description="The run_metadata to filter the step runs by.",
578
+ )
574
579
 
575
580
  model_config = ConfigDict(protected_namespaces=())
576
581
 
@@ -589,6 +594,7 @@ class StepRunFilter(WorkspaceScopedFilter):
589
594
  from zenml.zen_stores.schemas import (
590
595
  ModelSchema,
591
596
  ModelVersionSchema,
597
+ RunMetadataSchema,
592
598
  StepRunSchema,
593
599
  )
594
600
 
@@ -601,5 +607,21 @@ class StepRunFilter(WorkspaceScopedFilter):
601
607
  ),
602
608
  )
603
609
  custom_filters.append(model_filter)
610
+ if self.run_metadata is not None:
611
+ from zenml.enums import MetadataResourceTypes
612
+
613
+ for key, value in self.run_metadata.items():
614
+ additional_filter = and_(
615
+ RunMetadataSchema.resource_id == StepRunSchema.id,
616
+ RunMetadataSchema.resource_type
617
+ == MetadataResourceTypes.STEP_RUN,
618
+ RunMetadataSchema.key == key,
619
+ self.generate_custom_query_conditions_for_column(
620
+ value=value,
621
+ table=RunMetadataSchema,
622
+ column="value",
623
+ ),
624
+ )
625
+ custom_filters.append(additional_filter)
604
626
 
605
627
  return custom_filters
@@ -15,6 +15,7 @@
15
15
 
16
16
  from abc import ABC, abstractmethod
17
17
  from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional, Type, cast
18
+ from uuid import UUID
18
19
 
19
20
  from pydantic import model_validator
20
21
 
@@ -186,7 +187,17 @@ class BaseOrchestrator(StackComponent, ABC):
186
187
  """
187
188
  self._prepare_run(deployment=deployment)
188
189
 
189
- environment = get_config_environment_vars(deployment=deployment)
190
+ pipeline_run_id: Optional[UUID] = None
191
+ schedule_id: Optional[UUID] = None
192
+ if deployment.schedule:
193
+ schedule_id = deployment.schedule.id
194
+ if placeholder_run:
195
+ pipeline_run_id = placeholder_run.id
196
+
197
+ environment = get_config_environment_vars(
198
+ schedule_id=schedule_id,
199
+ pipeline_run_id=pipeline_run_id,
200
+ )
190
201
 
191
202
  prevent_client_side_caching = handle_bool_env_var(
192
203
  ENV_ZENML_PREVENT_CLIENT_SIDE_CACHING, default=False
@@ -422,7 +422,8 @@ class StepLauncher:
422
422
  )
423
423
  )
424
424
  environment = orchestrator_utils.get_config_environment_vars(
425
- deployment=self._deployment
425
+ pipeline_run_id=step_run_info.run_id,
426
+ step_run_id=step_run_info.step_run_id,
426
427
  )
427
428
  if last_retry:
428
429
  environment[ENV_ZENML_IGNORE_FAILURE_HOOK] = str(False)
@@ -28,16 +28,16 @@ from zenml.constants import (
28
28
  ENV_ZENML_DISABLE_CREDENTIALS_DISK_CACHING,
29
29
  ENV_ZENML_SERVER,
30
30
  ENV_ZENML_STORE_PREFIX,
31
- PIPELINE_API_TOKEN_EXPIRES_MINUTES,
32
31
  )
33
32
  from zenml.enums import AuthScheme, StackComponentType, StoreType
34
33
  from zenml.logger import get_logger
35
34
  from zenml.stack import StackComponent
36
35
  from zenml.utils.string_utils import format_name_template
37
36
 
37
+ logger = get_logger(__name__)
38
+
38
39
  if TYPE_CHECKING:
39
40
  from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
40
- from zenml.models import PipelineDeploymentResponse
41
41
 
42
42
 
43
43
  def get_orchestrator_run_name(pipeline_name: str) -> str:
@@ -80,16 +80,23 @@ def is_setting_enabled(
80
80
 
81
81
 
82
82
  def get_config_environment_vars(
83
- deployment: Optional["PipelineDeploymentResponse"] = None,
83
+ schedule_id: Optional[UUID] = None,
84
+ pipeline_run_id: Optional[UUID] = None,
85
+ step_run_id: Optional[UUID] = None,
84
86
  ) -> Dict[str, str]:
85
87
  """Gets environment variables to set for mirroring the active config.
86
88
 
87
- If a pipeline deployment is given, the environment variables will be set to
88
- include a newly generated API token valid for the duration of the pipeline
89
- run instead of the API token from the global config.
89
+ If a schedule ID, pipeline run ID or step run ID is given, and the current
90
+ client is not authenticated to a server with an API key, the environment
91
+ variables will be updated to include a newly generated workload API token
92
+ that will be valid for the duration of the schedule, pipeline run, or step
93
+ run instead of the current API token used to authenticate the client.
90
94
 
91
95
  Args:
92
- deployment: Optional deployment to use for the environment variables.
96
+ schedule_id: Optional schedule ID to use to generate a new API token.
97
+ pipeline_run_id: Optional pipeline run ID to use to generate a new API
98
+ token.
99
+ step_run_id: Optional step run ID to use to generate a new API token.
93
100
 
94
101
  Returns:
95
102
  Environment variable dict.
@@ -107,34 +114,46 @@ def get_config_environment_vars(
107
114
  ):
108
115
  credentials_store = get_credentials_store()
109
116
  url = global_config.store_configuration.url
110
- api_key = credentials_store.get_api_key(url)
111
117
  api_token = credentials_store.get_token(url, allow_expired=False)
112
- if api_key:
113
- environment_vars[ENV_ZENML_STORE_PREFIX + "API_KEY"] = api_key
114
- elif deployment:
115
- # When connected to an authenticated ZenML server, if a pipeline
116
- # deployment is supplied, we need to fetch an API token that will be
117
- # valid for the duration of the pipeline run.
118
+ if schedule_id or pipeline_run_id or step_run_id:
119
+ # When connected to an authenticated ZenML server, if a schedule ID,
120
+ # pipeline run ID or step run ID is supplied, we need to fetch a new
121
+ # workload API token scoped to the schedule, pipeline run or step
122
+ # run.
118
123
  assert isinstance(global_config.zen_store, RestZenStore)
119
- pipeline_id: Optional[UUID] = None
120
- if deployment.pipeline:
121
- pipeline_id = deployment.pipeline.id
122
- schedule_id: Optional[UUID] = None
123
- expires_minutes: Optional[int] = PIPELINE_API_TOKEN_EXPIRES_MINUTES
124
- if deployment.schedule:
125
- schedule_id = deployment.schedule.id
126
- # If a schedule is given, this is a long running pipeline that
127
- # should not have an API token that expires.
128
- expires_minutes = None
124
+
125
+ # If only a schedule is given, the pipeline run credentials will
126
+ # be valid for the entire duration of the schedule.
127
+ api_key = credentials_store.get_api_key(url)
128
+ if not api_key and not pipeline_run_id and not step_run_id:
129
+ logger.warning(
130
+ "An API token without an expiration time will be generated "
131
+ "and used to run this pipeline on a schedule. This is very "
132
+ "insecure because the API token will be valid for the "
133
+ "entire lifetime of the schedule and can be used to access "
134
+ "your user account if accidentally leaked. When deploying "
135
+ "a pipeline on a schedule, it is strongly advised to use a "
136
+ "service account API key to authenticate to the ZenML "
137
+ "server instead of your regular user account. For more "
138
+ "information, see "
139
+ "https://docs.zenml.io/how-to/connecting-to-zenml/connect-with-a-service-account"
140
+ )
141
+
142
+ # The schedule, pipeline run or step run credentials are scoped to
143
+ # the schedule, pipeline run or step run and will only be valid for
144
+ # the duration of the schedule/pipeline run/step run.
129
145
  new_api_token = global_config.zen_store.get_api_token(
130
- pipeline_id=pipeline_id,
131
146
  schedule_id=schedule_id,
132
- expires_minutes=expires_minutes,
147
+ pipeline_run_id=pipeline_run_id,
148
+ step_run_id=step_run_id,
133
149
  )
150
+
134
151
  environment_vars[ENV_ZENML_STORE_PREFIX + "API_TOKEN"] = (
135
152
  new_api_token
136
153
  )
137
154
  elif api_token:
155
+ # For all other cases, the pipeline run environment is configured
156
+ # with the current access token.
138
157
  environment_vars[ENV_ZENML_STORE_PREFIX + "API_TOKEN"] = (
139
158
  api_token.access_token
140
159
  )
zenml/steps/utils.py CHANGED
@@ -442,6 +442,11 @@ def log_step_metadata(
442
442
  from within a step or if no pipeline name or ID is provided and
443
443
  the function is not called from within a step.
444
444
  """
445
+ logger.warning(
446
+ "The `log_step_metadata` function is deprecated and will soon be "
447
+ "removed. Please use `log_metadata` instead."
448
+ )
449
+
445
450
  step_context = None
446
451
  if not step_name:
447
452
  with contextlib.suppress(RuntimeError):
@@ -0,0 +1,335 @@
1
+ # Copyright (c) ZenML GmbH 2024. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at:
6
+ #
7
+ # https://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
12
+ # or implied. See the License for the specific language governing
13
+ # permissions and limitations under the License.
14
+ """Utility functions to handle metadata for ZenML entities."""
15
+
16
+ import contextlib
17
+ from typing import Dict, Optional, Union, overload
18
+ from uuid import UUID
19
+
20
+ from zenml.client import Client
21
+ from zenml.enums import MetadataResourceTypes
22
+ from zenml.logger import get_logger
23
+ from zenml.metadata.metadata_types import MetadataType
24
+ from zenml.steps.step_context import get_step_context
25
+
26
+ logger = get_logger(__name__)
27
+
28
+
29
+ @overload
30
+ def log_metadata(metadata: Dict[str, MetadataType]) -> None: ...
31
+
32
+
33
+ @overload
34
+ def log_metadata(
35
+ *,
36
+ metadata: Dict[str, MetadataType],
37
+ artifact_version_id: UUID,
38
+ ) -> None: ...
39
+
40
+
41
+ @overload
42
+ def log_metadata(
43
+ *,
44
+ metadata: Dict[str, MetadataType],
45
+ artifact_name: str,
46
+ artifact_version: Optional[str] = None,
47
+ ) -> None: ...
48
+
49
+
50
+ @overload
51
+ def log_metadata(
52
+ *,
53
+ metadata: Dict[str, MetadataType],
54
+ model_version_id: UUID,
55
+ ) -> None: ...
56
+
57
+
58
+ @overload
59
+ def log_metadata(
60
+ *,
61
+ metadata: Dict[str, MetadataType],
62
+ model_name: str,
63
+ model_version: str,
64
+ ) -> None: ...
65
+
66
+
67
+ @overload
68
+ def log_metadata(
69
+ *,
70
+ metadata: Dict[str, MetadataType],
71
+ step_id: UUID,
72
+ ) -> None: ...
73
+
74
+
75
+ @overload
76
+ def log_metadata(
77
+ *,
78
+ metadata: Dict[str, MetadataType],
79
+ run_id_name_or_prefix: Union[UUID, str],
80
+ ) -> None: ...
81
+
82
+
83
+ @overload
84
+ def log_metadata(
85
+ *,
86
+ metadata: Dict[str, MetadataType],
87
+ step_name: str,
88
+ run_id_name_or_prefix: Union[UUID, str],
89
+ ) -> None: ...
90
+
91
+
92
+ def log_metadata(
93
+ metadata: Dict[str, MetadataType],
94
+ # Parameters to manually log metadata for steps and runs
95
+ step_id: Optional[UUID] = None,
96
+ step_name: Optional[str] = None,
97
+ run_id_name_or_prefix: Optional[Union[UUID, str]] = None,
98
+ # Parameters to manually log metadata for artifacts
99
+ artifact_version_id: Optional[UUID] = None,
100
+ artifact_name: Optional[str] = None,
101
+ artifact_version: Optional[str] = None,
102
+ # Parameters to manually log metadata for models
103
+ model_version_id: Optional[UUID] = None,
104
+ model_name: Optional[str] = None,
105
+ model_version: Optional[str] = None,
106
+ ) -> None:
107
+ """Logs metadata for various resource types in a generalized way.
108
+
109
+ Args:
110
+ metadata: The metadata to log.
111
+ step_id: The ID of the step.
112
+ step_name: The name of the step.
113
+ run_id_name_or_prefix: The id, name or prefix of the run
114
+ artifact_version_id: The ID of the artifact version
115
+ artifact_name: The name of the artifact.
116
+ artifact_version: The version of the artifact.
117
+ model_version_id: The ID of the model version.
118
+ model_name: The name of the model.
119
+ model_version: The version of the model
120
+
121
+ Raises:
122
+ ValueError: If no identifiers are provided and the function is not
123
+ called from within a step.
124
+ """
125
+ client = Client()
126
+
127
+ # If a step name is provided, we need a run_id_name_or_prefix and will log
128
+ # metadata for the steps pipeline and model accordingly.
129
+ if step_name is not None and run_id_name_or_prefix is not None:
130
+ run_model = client.get_pipeline_run(
131
+ name_id_or_prefix=run_id_name_or_prefix
132
+ )
133
+ step_model = run_model.steps[step_name]
134
+
135
+ client.create_run_metadata(
136
+ metadata=metadata,
137
+ resource_id=run_model.id,
138
+ resource_type=MetadataResourceTypes.PIPELINE_RUN,
139
+ )
140
+ client.create_run_metadata(
141
+ metadata=metadata,
142
+ resource_id=step_model.id,
143
+ resource_type=MetadataResourceTypes.STEP_RUN,
144
+ )
145
+ if step_model.model_version:
146
+ client.create_run_metadata(
147
+ metadata=metadata,
148
+ resource_id=step_model.model_version.id,
149
+ resource_type=MetadataResourceTypes.MODEL_VERSION,
150
+ )
151
+
152
+ # If a step is identified by id, fetch it directly through the client,
153
+ # follow a similar procedure and log metadata for its pipeline and model
154
+ # as well.
155
+ elif step_id is not None:
156
+ step_model = client.get_run_step(step_run_id=step_id)
157
+ client.create_run_metadata(
158
+ metadata=metadata,
159
+ resource_id=step_model.pipeline_run_id,
160
+ resource_type=MetadataResourceTypes.PIPELINE_RUN,
161
+ )
162
+ client.create_run_metadata(
163
+ metadata=metadata,
164
+ resource_id=step_model.id,
165
+ resource_type=MetadataResourceTypes.STEP_RUN,
166
+ )
167
+ if step_model.model_version:
168
+ client.create_run_metadata(
169
+ metadata=metadata,
170
+ resource_id=step_model.model_version.id,
171
+ resource_type=MetadataResourceTypes.MODEL_VERSION,
172
+ )
173
+
174
+ # If a pipeline run id is identified, we need to log metadata to it and its
175
+ # model as well.
176
+ elif run_id_name_or_prefix is not None:
177
+ run_model = client.get_pipeline_run(
178
+ name_id_or_prefix=run_id_name_or_prefix
179
+ )
180
+ client.create_run_metadata(
181
+ metadata=metadata,
182
+ resource_id=run_model.id,
183
+ resource_type=MetadataResourceTypes.PIPELINE_RUN,
184
+ )
185
+ if run_model.model_version:
186
+ client.create_run_metadata(
187
+ metadata=metadata,
188
+ resource_id=run_model.model_version.id,
189
+ resource_type=MetadataResourceTypes.MODEL_VERSION,
190
+ )
191
+
192
+ # If the user provides a model name and version, we use to model abstraction
193
+ # to fetch the model version and attach the corresponding metadata to it.
194
+ elif model_name is not None and model_version is not None:
195
+ from zenml import Model
196
+
197
+ mv = Model(name=model_name, version=model_version)
198
+ client.create_run_metadata(
199
+ metadata=metadata,
200
+ resource_id=mv.id,
201
+ resource_type=MetadataResourceTypes.MODEL_VERSION,
202
+ )
203
+
204
+ # If the user provides a model version id, we use the client to fetch it and
205
+ # attach the metadata to it.
206
+ elif model_version_id is not None:
207
+ model_version_id = client.get_model_version(
208
+ model_version_name_or_number_or_id=model_version_id
209
+ ).id
210
+ client.create_run_metadata(
211
+ metadata=metadata,
212
+ resource_id=model_version_id,
213
+ resource_type=MetadataResourceTypes.MODEL_VERSION,
214
+ )
215
+
216
+ # If the user provides an artifact name, there are three possibilities. If
217
+ # an artifact version is also provided with the name, we use both to fetch
218
+ # the artifact version and use it to log the metadata. If no version is
219
+ # provided, if the function is called within a step we search the artifacts
220
+ # of the step if not we fetch the latest version and attach the metadata
221
+ # to the latest version.
222
+ elif artifact_name is not None:
223
+ if artifact_version:
224
+ artifact_version_model = client.get_artifact_version(
225
+ name_id_or_prefix=artifact_name, version=artifact_version
226
+ )
227
+ client.create_run_metadata(
228
+ metadata=metadata,
229
+ resource_id=artifact_version_model.id,
230
+ resource_type=MetadataResourceTypes.ARTIFACT_VERSION,
231
+ )
232
+ else:
233
+ step_context = None
234
+ with contextlib.suppress(RuntimeError):
235
+ step_context = get_step_context()
236
+
237
+ if step_context:
238
+ step_context.add_output_metadata(
239
+ metadata=metadata, output_name=artifact_name
240
+ )
241
+ else:
242
+ artifact_version_model = client.get_artifact_version(
243
+ name_id_or_prefix=artifact_name
244
+ )
245
+ client.create_run_metadata(
246
+ metadata=metadata,
247
+ resource_id=artifact_version_model.id,
248
+ resource_type=MetadataResourceTypes.ARTIFACT_VERSION,
249
+ )
250
+
251
+ # If the user directly provides an artifact_version_id, we use the client to
252
+ # fetch is and attach the metadata accordingly.
253
+ elif artifact_version_id is not None:
254
+ artifact_version_model = client.get_artifact_version(
255
+ name_id_or_prefix=artifact_version_id,
256
+ )
257
+ client.create_run_metadata(
258
+ metadata=metadata,
259
+ resource_id=artifact_version_model.id,
260
+ resource_type=MetadataResourceTypes.ARTIFACT_VERSION,
261
+ )
262
+
263
+ # If every additional value is None, that means we are calling it bare bones
264
+ # and this call needs to happen during a step execution. We will use the
265
+ # step context to fetch the step, run and possibly the model version and
266
+ # attach the metadata accordingly.
267
+ elif all(
268
+ v is None
269
+ for v in [
270
+ step_id,
271
+ step_name,
272
+ run_id_name_or_prefix,
273
+ artifact_version_id,
274
+ artifact_name,
275
+ artifact_version,
276
+ model_version_id,
277
+ model_name,
278
+ model_version,
279
+ ]
280
+ ):
281
+ try:
282
+ step_context = get_step_context()
283
+ except RuntimeError:
284
+ raise ValueError(
285
+ "You are calling 'log_metadata()' outside of a step execution. "
286
+ "If you would like to add metadata to a ZenML entity outside "
287
+ "of the step execution, please provide the required "
288
+ "identifiers."
289
+ )
290
+ client.create_run_metadata(
291
+ metadata=metadata,
292
+ resource_id=step_context.pipeline_run.id,
293
+ resource_type=MetadataResourceTypes.PIPELINE_RUN,
294
+ )
295
+ client.create_run_metadata(
296
+ metadata=metadata,
297
+ resource_id=step_context.step_run.id,
298
+ resource_type=MetadataResourceTypes.STEP_RUN,
299
+ )
300
+ if step_context.model_version:
301
+ client.create_run_metadata(
302
+ metadata=metadata,
303
+ resource_id=step_context.model_version.id,
304
+ resource_type=MetadataResourceTypes.MODEL_VERSION,
305
+ )
306
+
307
+ else:
308
+ raise ValueError(
309
+ """
310
+ Unsupported way to call the `log_metadata`. Possible combinations "
311
+ include:
312
+
313
+ # Inside a step
314
+ # Logs the metadata to the step, its run and possibly its model
315
+ log_metadata(metadata={})
316
+
317
+ # Manually logging for a step
318
+ # Logs the metadata to the step, its run and possibly its model
319
+ log_metadata(metadata={}, step_name=..., run_id_name_or_prefix=...)
320
+ log_metadata(metadata={}, step_id=...)
321
+
322
+ # Manually logging for a run
323
+ # Logs the metadata to the run, possibly its model
324
+ log_metadata(metadata={}, run_id_name_or_prefix=...)
325
+
326
+ # Manually logging for a model
327
+ log_metadata(metadata={}, model_name=..., model_version=...)
328
+ log_metadata(metadata={}, model_version_id=...)
329
+
330
+ # Manually logging for an artifact
331
+ log_metadata(metadata={}, artifact_name=...) # inside a step
332
+ log_metadata(metadata={}, artifact_name=..., artifact_version=...)
333
+ log_metadata(metadata={}, artifact_version_id=...)
334
+ """
335
+ )