zenml-nightly 0.64.0.dev20240811__py3-none-any.whl → 0.66.0.dev20240910__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 (291) hide show
  1. README.md +1 -1
  2. RELEASE_NOTES.md +126 -4
  3. zenml/VERSION +1 -1
  4. zenml/artifacts/utils.py +13 -6
  5. zenml/cli/__init__.py +1 -1
  6. zenml/cli/base.py +4 -4
  7. zenml/cli/integration.py +48 -9
  8. zenml/cli/pipeline.py +9 -2
  9. zenml/cli/stack.py +39 -27
  10. zenml/cli/utils.py +13 -0
  11. zenml/client.py +15 -17
  12. zenml/config/compiler.py +34 -0
  13. zenml/config/server_config.py +30 -0
  14. zenml/config/source.py +3 -7
  15. zenml/constants.py +5 -3
  16. zenml/entrypoints/base_entrypoint_configuration.py +41 -27
  17. zenml/entrypoints/step_entrypoint_configuration.py +5 -2
  18. zenml/enums.py +2 -0
  19. zenml/environment.py +31 -0
  20. zenml/feature_stores/base_feature_store.py +4 -6
  21. zenml/integrations/__init__.py +3 -0
  22. zenml/integrations/airflow/flavors/airflow_orchestrator_flavor.py +9 -0
  23. zenml/integrations/aws/__init__.py +2 -2
  24. zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py +2 -1
  25. zenml/integrations/azure/__init__.py +2 -2
  26. zenml/integrations/azure/azureml_utils.py +201 -0
  27. zenml/integrations/azure/flavors/azureml.py +139 -0
  28. zenml/integrations/azure/flavors/azureml_orchestrator_flavor.py +20 -118
  29. zenml/integrations/azure/flavors/azureml_step_operator_flavor.py +67 -14
  30. zenml/integrations/azure/orchestrators/azureml_orchestrator.py +58 -172
  31. zenml/integrations/azure/orchestrators/azureml_orchestrator_entrypoint_config.py +1 -0
  32. zenml/integrations/azure/service_connectors/azure_service_connector.py +4 -0
  33. zenml/integrations/azure/step_operators/azureml_step_operator.py +78 -177
  34. zenml/integrations/constants.py +3 -0
  35. zenml/integrations/databricks/__init__.py +22 -4
  36. zenml/integrations/databricks/flavors/databricks_orchestrator_flavor.py +9 -0
  37. zenml/integrations/deepchecks/__init__.py +29 -11
  38. zenml/integrations/deepchecks/materializers/deepchecks_dataset_materializer.py +3 -1
  39. zenml/integrations/deepchecks/validation_checks.py +0 -30
  40. zenml/integrations/evidently/__init__.py +17 -2
  41. zenml/integrations/facets/__init__.py +21 -5
  42. zenml/integrations/feast/__init__.py +19 -6
  43. zenml/integrations/gcp/__init__.py +2 -2
  44. zenml/integrations/gcp/flavors/vertex_orchestrator_flavor.py +9 -0
  45. zenml/integrations/gcp/orchestrators/vertex_orchestrator.py +10 -1
  46. zenml/integrations/great_expectations/__init__.py +21 -7
  47. zenml/integrations/huggingface/__init__.py +39 -15
  48. zenml/integrations/huggingface/materializers/__init__.py +3 -0
  49. zenml/integrations/huggingface/materializers/huggingface_datasets_materializer.py +3 -1
  50. zenml/integrations/huggingface/materializers/huggingface_pt_model_materializer.py +1 -1
  51. zenml/integrations/huggingface/materializers/huggingface_t5_materializer.py +107 -0
  52. zenml/integrations/huggingface/materializers/huggingface_tf_model_materializer.py +1 -1
  53. zenml/integrations/huggingface/materializers/huggingface_tokenizer_materializer.py +2 -2
  54. zenml/integrations/huggingface/steps/accelerate_runner.py +108 -85
  55. zenml/integrations/hyperai/flavors/hyperai_orchestrator_flavor.py +9 -0
  56. zenml/integrations/kubeflow/flavors/kubeflow_orchestrator_flavor.py +9 -0
  57. zenml/integrations/kubeflow/orchestrators/kubeflow_orchestrator.py +10 -1
  58. zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py +9 -0
  59. zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py +10 -1
  60. zenml/integrations/lightning/__init__.py +48 -0
  61. zenml/integrations/lightning/flavors/__init__.py +23 -0
  62. zenml/integrations/lightning/flavors/lightning_orchestrator_flavor.py +148 -0
  63. zenml/integrations/lightning/orchestrators/__init__.py +23 -0
  64. zenml/integrations/lightning/orchestrators/lightning_orchestrator.py +596 -0
  65. zenml/integrations/lightning/orchestrators/lightning_orchestrator_entrypoint.py +307 -0
  66. zenml/integrations/lightning/orchestrators/lightning_orchestrator_entrypoint_configuration.py +77 -0
  67. zenml/integrations/lightning/orchestrators/utils.py +67 -0
  68. zenml/integrations/mlflow/__init__.py +43 -5
  69. zenml/integrations/mlflow/services/mlflow_deployment.py +26 -0
  70. zenml/integrations/numpy/__init__.py +32 -0
  71. zenml/integrations/numpy/materializers/__init__.py +18 -0
  72. zenml/integrations/numpy/materializers/numpy_materializer.py +246 -0
  73. zenml/integrations/pandas/__init__.py +32 -0
  74. zenml/integrations/pandas/materializers/__init__.py +18 -0
  75. zenml/integrations/pandas/materializers/pandas_materializer.py +192 -0
  76. zenml/integrations/prodigy/annotators/prodigy_annotator.py +1 -1
  77. zenml/integrations/seldon/__init__.py +18 -3
  78. zenml/integrations/sklearn/__init__.py +1 -1
  79. zenml/integrations/skypilot_azure/__init__.py +1 -1
  80. zenml/integrations/tensorboard/__init__.py +1 -1
  81. zenml/integrations/tensorflow/__init__.py +2 -2
  82. zenml/integrations/tensorflow/materializers/tf_dataset_materializer.py +2 -2
  83. zenml/integrations/whylogs/__init__.py +18 -2
  84. zenml/logging/step_logging.py +9 -2
  85. zenml/materializers/__init__.py +0 -4
  86. zenml/materializers/base_materializer.py +4 -0
  87. zenml/materializers/numpy_materializer.py +23 -234
  88. zenml/materializers/pandas_materializer.py +22 -179
  89. zenml/model/model.py +91 -2
  90. zenml/model/utils.py +5 -5
  91. zenml/models/__init__.py +16 -3
  92. zenml/models/v2/core/model_version.py +1 -1
  93. zenml/models/v2/core/pipeline_run.py +31 -1
  94. zenml/models/v2/core/stack.py +51 -20
  95. zenml/models/v2/core/step_run.py +28 -0
  96. zenml/models/v2/misc/info_models.py +78 -0
  97. zenml/new/pipelines/pipeline.py +65 -25
  98. zenml/new/pipelines/run_utils.py +57 -136
  99. zenml/new/steps/step_context.py +17 -6
  100. zenml/orchestrators/base_orchestrator.py +9 -0
  101. zenml/orchestrators/step_launcher.py +37 -14
  102. zenml/orchestrators/step_runner.py +14 -13
  103. zenml/orchestrators/utils.py +107 -7
  104. zenml/service_connectors/service_connector_utils.py +2 -2
  105. zenml/stack/utils.py +11 -2
  106. zenml/stack_deployments/azure_stack_deployment.py +2 -1
  107. zenml/steps/base_step.py +62 -25
  108. zenml/steps/utils.py +115 -3
  109. zenml/utils/cloud_utils.py +8 -8
  110. zenml/utils/code_utils.py +130 -32
  111. zenml/utils/function_utils.py +7 -7
  112. zenml/utils/notebook_utils.py +14 -0
  113. zenml/utils/pipeline_docker_image_builder.py +1 -11
  114. zenml/utils/pydantic_utils.py +3 -3
  115. zenml/utils/secret_utils.py +2 -2
  116. zenml/utils/settings_utils.py +1 -1
  117. zenml/utils/source_utils.py +67 -21
  118. zenml/utils/string_utils.py +29 -0
  119. zenml/zen_server/dashboard/assets/{404-CRAA_Lew.js → 404-iO8vpun1.js} +1 -1
  120. zenml/zen_server/dashboard/assets/{@radix-BXWm7HOa.js → @radix-DnFH_oo1.js} +1 -1
  121. zenml/zen_server/dashboard/assets/{@react-router-l3lMcXA2.js → @react-router-APVeuk-U.js} +1 -1
  122. zenml/zen_server/dashboard/assets/{@reactflow-CeVxyqYT.js → @reactflow-B6kq9fJZ.js} +2 -2
  123. zenml/zen_server/dashboard/assets/{@tanstack-FmcYZMuX.js → @tanstack-QbMbTrh5.js} +1 -1
  124. zenml/zen_server/dashboard/assets/AlertDialogDropdownItem-BXeSvmMY.js +1 -0
  125. zenml/zen_server/dashboard/assets/{CodeSnippet-D0VLxT2A.js → CodeSnippet-DNWdQmbo.js} +2 -2
  126. zenml/zen_server/dashboard/assets/CollapsibleCard-B2OVjWYE.js +1 -0
  127. zenml/zen_server/dashboard/assets/Commands-DsoaVElZ.js +1 -0
  128. zenml/zen_server/dashboard/assets/CopyButton-BqE_-PHO.js +2 -0
  129. zenml/zen_server/dashboard/assets/{CsvVizualization-D3kAypDj.js → CsvVizualization-Dyasr2jU.js} +6 -6
  130. zenml/zen_server/dashboard/assets/{edit-C0MVvPD2.js → DialogItem-Cz1VLRwa.js} +1 -1
  131. zenml/zen_server/dashboard/assets/{DisplayDate-DizbSeT-.js → DisplayDate-DkCy54Bp.js} +1 -1
  132. zenml/zen_server/dashboard/assets/EditSecretDialog-Du423_3U.js +1 -0
  133. zenml/zen_server/dashboard/assets/{EmptyState-BHblM39I.js → EmptyState-Cs3DEmso.js} +1 -1
  134. zenml/zen_server/dashboard/assets/{Error-C6LeJSER.js → Error-DorJD_va.js} +1 -1
  135. zenml/zen_server/dashboard/assets/ExecutionStatus-CIfQTutR.js +1 -0
  136. zenml/zen_server/dashboard/assets/{Helpbox-aAB2XP-z.js → Helpbox-CmfvtNeq.js} +1 -1
  137. zenml/zen_server/dashboard/assets/Infobox-BL9NOS37.js +1 -0
  138. zenml/zen_server/dashboard/assets/{InlineAvatar-DpTLgM3Q.js → InlineAvatar-Ds2ZFHPc.js} +1 -1
  139. zenml/zen_server/dashboard/assets/{Lock-CNyJvf2r.js → Lock-CmIn0szs.js} +1 -1
  140. zenml/zen_server/dashboard/assets/{MarkdownVisualization-Bajxn0HY.js → MarkdownVisualization-DS05sfBm.js} +1 -1
  141. zenml/zen_server/dashboard/assets/{NumberBox-BmKE0qnO.js → NumberBox-CrN0_kqI.js} +1 -1
  142. zenml/zen_server/dashboard/assets/Partials-DX-8iEa1.js +1 -0
  143. zenml/zen_server/dashboard/assets/{PasswordChecker-yGGoJSB-.js → PasswordChecker-DE71J_3F.js} +1 -1
  144. zenml/zen_server/dashboard/assets/ProviderIcon-BOQJgapd.js +1 -0
  145. zenml/zen_server/dashboard/assets/ProviderRadio-BsYBw9YA.js +1 -0
  146. zenml/zen_server/dashboard/assets/SearchField-W3GXpLlI.js +1 -0
  147. zenml/zen_server/dashboard/assets/SetPassword-B-0a8UCj.js +1 -0
  148. zenml/zen_server/dashboard/assets/{Tick-uxv80Q6a.js → Tick-i1DYsVcX.js} +1 -1
  149. zenml/zen_server/dashboard/assets/{UpdatePasswordSchemas-oN4G3sKz.js → UpdatePasswordSchemas-C6Zb7ASL.js} +1 -1
  150. zenml/zen_server/dashboard/assets/UsageReason-CCnzmwS8.js +1 -0
  151. zenml/zen_server/dashboard/assets/WizardFooter-BHbO7zOa.js +1 -0
  152. zenml/zen_server/dashboard/assets/all-pipeline-runs-query-BBEe6I9-.js +1 -0
  153. zenml/zen_server/dashboard/assets/{check-circle-1_I207rW.js → check-circle-DOoS4yhF.js} +1 -1
  154. zenml/zen_server/dashboard/assets/{chevron-down-BpaF8JqM.js → chevron-down-Cwb-W_B_.js} +1 -1
  155. zenml/zen_server/dashboard/assets/{chevron-right-double-Dk8e2L99.js → chevron-right-double-c9H46Kl8.js} +1 -1
  156. zenml/zen_server/dashboard/assets/{cloud-only-BkUuI0lZ.js → cloud-only-BuP4Kt_7.js} +1 -1
  157. zenml/zen_server/dashboard/assets/code-browser-BJYErIjr.js +1 -0
  158. zenml/zen_server/dashboard/assets/codespaces-BitYDX9d.gif +0 -0
  159. zenml/zen_server/dashboard/assets/{copy-f3XGPPxt.js → copy-CaGlDsUy.js} +1 -1
  160. zenml/zen_server/dashboard/assets/create-stack-B2x2d4r1.js +1 -0
  161. zenml/zen_server/dashboard/assets/{docker-8uj__HHK.js → docker-BFAFXr2_.js} +1 -1
  162. zenml/zen_server/dashboard/assets/{dots-horizontal-sKQlWEni.js → dots-horizontal-C6K59vUm.js} +1 -1
  163. zenml/zen_server/dashboard/assets/flyte-Cj-xy_8I.svg +10 -0
  164. zenml/zen_server/dashboard/assets/form-schemas-Bap0f854.js +1 -0
  165. zenml/zen_server/dashboard/assets/gcp-Dj6ntk0L.js +1 -0
  166. zenml/zen_server/dashboard/assets/{help-FuHlZwn0.js → help-CwN931fX.js} +1 -1
  167. zenml/zen_server/dashboard/assets/{index-Bd1xgUQG.js → index-5GJ5ysEZ.js} +1 -1
  168. zenml/zen_server/dashboard/assets/{index-DaGknux4.css → index-6DYjZgDn.css} +1 -1
  169. zenml/zen_server/dashboard/assets/index-B9wVwe7u.js +55 -0
  170. zenml/zen_server/dashboard/assets/index-DFi8BroH.js +1 -0
  171. zenml/zen_server/dashboard/assets/{index.esm-DT4uyn2i.js → index.esm-BE1uqCX5.js} +1 -1
  172. zenml/zen_server/dashboard/assets/kubernetes-BjbR6D-1.js +1 -0
  173. zenml/zen_server/dashboard/assets/{layout-D6oiSbfd.js → layout-Dru15_XR.js} +1 -1
  174. zenml/zen_server/dashboard/assets/link-external-BT2L8hAQ.js +1 -0
  175. zenml/zen_server/dashboard/assets/{login-mutation-13A_JSVA.js → login-mutation-DwxUz8VA.js} +1 -1
  176. zenml/zen_server/dashboard/assets/{logs-CgeE2vZP.js → logs-GiDJXbLS.js} +1 -1
  177. zenml/zen_server/dashboard/assets/metaflow-weOkWNyT.svg +10 -0
  178. zenml/zen_server/dashboard/assets/{not-found-B0Mmb90p.js → not-found-D5i9DunU.js} +1 -1
  179. zenml/zen_server/dashboard/assets/{package-DdkziX79.js → package-DYKZ5jKW.js} +1 -1
  180. zenml/zen_server/dashboard/assets/page-BFuJICXM.js +9 -0
  181. zenml/zen_server/dashboard/assets/{page-BGwA9B1M.js → page-BiF8hLbO.js} +1 -1
  182. zenml/zen_server/dashboard/assets/{page-DugsjcQ_.js → page-BitfWsiW.js} +1 -1
  183. zenml/zen_server/dashboard/assets/page-CDOQLrPC.js +1 -0
  184. zenml/zen_server/dashboard/assets/page-CEJWu1YO.js +1 -0
  185. zenml/zen_server/dashboard/assets/page-CIbehp7V.js +1 -0
  186. zenml/zen_server/dashboard/assets/page-CLiRGfWo.js +1 -0
  187. zenml/zen_server/dashboard/assets/page-CV44mQn9.js +1 -0
  188. zenml/zen_server/dashboard/assets/page-CrSdkteO.js +2 -0
  189. zenml/zen_server/dashboard/assets/page-D5F3DJjm.js +1 -0
  190. zenml/zen_server/dashboard/assets/page-DE03uZZR.js +1 -0
  191. zenml/zen_server/dashboard/assets/page-DFCK65G9.js +1 -0
  192. zenml/zen_server/dashboard/assets/{page-RnG-qhv9.js → page-DGMa3ZQL.js} +1 -1
  193. zenml/zen_server/dashboard/assets/page-DI-qTWrm.js +1 -0
  194. zenml/zen_server/dashboard/assets/page-DQGCHKrQ.js +1 -0
  195. zenml/zen_server/dashboard/assets/{page-DSTQnBk-.js → page-DQdwZZ9x.js} +1 -1
  196. zenml/zen_server/dashboard/assets/page-DgM-N9RL.js +1 -0
  197. zenml/zen_server/dashboard/assets/page-Dt8VgzbE.js +1 -0
  198. zenml/zen_server/dashboard/assets/{page-DLpOnf7u.js → page-J0s8Sq3N.js} +1 -1
  199. zenml/zen_server/dashboard/assets/page-WCQ659by.js +1 -0
  200. zenml/zen_server/dashboard/assets/page-bimkItOg.js +1 -0
  201. zenml/zen_server/dashboard/assets/{page-hQaiQXfg.js → page-iwoJnwPv.js} +1 -1
  202. zenml/zen_server/dashboard/assets/{page-YiF_fNbe.js → page-oS4hqS8M.js} +1 -1
  203. zenml/zen_server/dashboard/assets/page-oSqx9dkH.js +1 -0
  204. zenml/zen_server/dashboard/assets/page-p3GqEAUW.js +1 -0
  205. zenml/zen_server/dashboard/assets/page-qvcUVPE-.js +1 -0
  206. zenml/zen_server/dashboard/assets/page-xQG6GmFJ.js +1 -0
  207. zenml/zen_server/dashboard/assets/{persist-3-5nOJ6m.js → persist-mEZN_fgH.js} +1 -1
  208. zenml/zen_server/dashboard/assets/persist-xsYgVtR1.js +1 -0
  209. zenml/zen_server/dashboard/assets/{plus-FB9-lEq_.js → plus-Bc8eLSDM.js} +1 -1
  210. zenml/zen_server/dashboard/assets/{refresh-COb6KYDi.js → refresh-hfgWPeto.js} +1 -1
  211. zenml/zen_server/dashboard/assets/rocket-SESCGQXm.js +1 -0
  212. zenml/zen_server/dashboard/assets/sharedSchema-BfZcy7aP.js +14 -0
  213. zenml/zen_server/dashboard/assets/stack-detail-query-CU4egfhp.js +1 -0
  214. zenml/zen_server/dashboard/assets/templates-1S_8WeSK.webp +0 -0
  215. zenml/zen_server/dashboard/assets/{trash-Cd5CSFqA.js → trash-DUWZWzse.js} +1 -1
  216. zenml/zen_server/dashboard/assets/{update-server-settings-mutation-B8GB_ubU.js → update-server-settings-mutation-DNqmQXDM.js} +1 -1
  217. zenml/zen_server/dashboard/assets/{url-hcMJkz8p.js → url-DwbuKk1b.js} +1 -1
  218. zenml/zen_server/dashboard/assets/{zod-CnykDKJj.js → zod-uFd1wBcd.js} +1 -1
  219. zenml/zen_server/dashboard/index.html +7 -7
  220. zenml/zen_server/dashboard_legacy/asset-manifest.json +4 -4
  221. zenml/zen_server/dashboard_legacy/index.html +1 -1
  222. zenml/zen_server/dashboard_legacy/{precache-manifest.9c473c96a43298343a7ce1256183123b.js → precache-manifest.290b95d5b43efa3368b3dc63d20c4782.js} +4 -4
  223. zenml/zen_server/dashboard_legacy/service-worker.js +1 -1
  224. zenml/zen_server/dashboard_legacy/static/js/{main.463c90b9.chunk.js → main.840d1bf0.chunk.js} +2 -2
  225. zenml/zen_server/dashboard_legacy/static/js/{main.463c90b9.chunk.js.map → main.840d1bf0.chunk.js.map} +1 -1
  226. zenml/zen_server/deploy/helm/Chart.yaml +1 -1
  227. zenml/zen_server/deploy/helm/README.md +2 -2
  228. zenml/zen_server/routers/service_connectors_endpoints.py +2 -4
  229. zenml/zen_server/routers/workspaces_endpoints.py +20 -66
  230. zenml/zen_server/secure_headers.py +120 -0
  231. zenml/zen_server/template_execution/runner_entrypoint_configuration.py +0 -2
  232. zenml/zen_server/template_execution/utils.py +1 -0
  233. zenml/zen_server/utils.py +0 -100
  234. zenml/zen_server/zen_server_api.py +4 -2
  235. zenml/zen_stores/migrations/versions/0.65.0_release.py +23 -0
  236. zenml/zen_stores/migrations/versions/0.66.0_release.py +23 -0
  237. zenml/zen_stores/migrations/versions/bf2120261b5a_add_configured_model_version_id.py +74 -0
  238. zenml/zen_stores/rest_zen_store.py +4 -21
  239. zenml/zen_stores/schemas/constants.py +16 -0
  240. zenml/zen_stores/schemas/model_schemas.py +9 -3
  241. zenml/zen_stores/schemas/pipeline_run_schemas.py +22 -8
  242. zenml/zen_stores/schemas/step_run_schemas.py +23 -12
  243. zenml/zen_stores/sql_zen_store.py +312 -300
  244. zenml/zen_stores/zen_store_interface.py +0 -16
  245. {zenml_nightly-0.64.0.dev20240811.dist-info → zenml_nightly-0.66.0.dev20240910.dist-info}/METADATA +10 -12
  246. {zenml_nightly-0.64.0.dev20240811.dist-info → zenml_nightly-0.66.0.dev20240910.dist-info}/RECORD +249 -217
  247. zenml/models/v2/misc/full_stack.py +0 -129
  248. zenml/new/pipelines/model_utils.py +0 -72
  249. zenml/zen_server/dashboard/assets/AlertDialogDropdownItem-ErO9aOgK.js +0 -1
  250. zenml/zen_server/dashboard/assets/AwarenessChannel-CLXo5rKM.js +0 -1
  251. zenml/zen_server/dashboard/assets/CollapsibleCard-BaUPiVg0.js +0 -1
  252. zenml/zen_server/dashboard/assets/Commands-JrcZK-3j.js +0 -1
  253. zenml/zen_server/dashboard/assets/CopyButton-Dbo52T1K.js +0 -2
  254. zenml/zen_server/dashboard/assets/EditSecretDialog-Bd7mFLS4.js +0 -1
  255. zenml/zen_server/dashboard/assets/ExecutionStatus-jH4OrWBq.js +0 -1
  256. zenml/zen_server/dashboard/assets/Infobox-BQ0aty32.js +0 -1
  257. zenml/zen_server/dashboard/assets/ProviderRadio-BBqkIuTd.js +0 -1
  258. zenml/zen_server/dashboard/assets/RadioItem-xLhXoiFV.js +0 -1
  259. zenml/zen_server/dashboard/assets/SearchField-C9R0mdaX.js +0 -1
  260. zenml/zen_server/dashboard/assets/SetPassword-52sNxNiO.js +0 -1
  261. zenml/zen_server/dashboard/assets/SuccessStep-DlkItqYG.js +0 -1
  262. zenml/zen_server/dashboard/assets/aws-0_3UsPif.js +0 -1
  263. zenml/zen_server/dashboard/assets/database-cXYNX9tt.js +0 -1
  264. zenml/zen_server/dashboard/assets/file-text-B9JibxTs.js +0 -1
  265. zenml/zen_server/dashboard/assets/index-DhIZtpxB.js +0 -55
  266. zenml/zen_server/dashboard/assets/page-7-v2OBm-.js +0 -1
  267. zenml/zen_server/dashboard/assets/page-B3ozwdD1.js +0 -1
  268. zenml/zen_server/dashboard/assets/page-BkjAUyTA.js +0 -1
  269. zenml/zen_server/dashboard/assets/page-BnacgBiy.js +0 -1
  270. zenml/zen_server/dashboard/assets/page-BxF_KMQ3.js +0 -2
  271. zenml/zen_server/dashboard/assets/page-C4POHC0K.js +0 -1
  272. zenml/zen_server/dashboard/assets/page-C9kudd44.js +0 -9
  273. zenml/zen_server/dashboard/assets/page-CA1j3GpJ.js +0 -1
  274. zenml/zen_server/dashboard/assets/page-CCY6yfmu.js +0 -1
  275. zenml/zen_server/dashboard/assets/page-CgTe7Bme.js +0 -1
  276. zenml/zen_server/dashboard/assets/page-Cgn-6v2Y.js +0 -1
  277. zenml/zen_server/dashboard/assets/page-CxQmQqDw.js +0 -1
  278. zenml/zen_server/dashboard/assets/page-D2Goey3H.js +0 -1
  279. zenml/zen_server/dashboard/assets/page-DTysUGOy.js +0 -1
  280. zenml/zen_server/dashboard/assets/page-D_EXUFJb.js +0 -1
  281. zenml/zen_server/dashboard/assets/page-Db15QzsM.js +0 -1
  282. zenml/zen_server/dashboard/assets/page-OFKSPyN7.js +0 -1
  283. zenml/zen_server/dashboard/assets/page-T2BtjwPl.js +0 -1
  284. zenml/zen_server/dashboard/assets/page-TXe1Eo3Z.js +0 -1
  285. zenml/zen_server/dashboard/assets/play-circle-XSkLR12B.js +0 -1
  286. zenml/zen_server/dashboard/assets/sharedSchema-BoYx_B_L.js +0 -14
  287. zenml/zen_server/dashboard/assets/stack-detail-query-B-US_-wa.js +0 -1
  288. zenml/zen_server/dashboard/assets/terminal-grtjrIEJ.js +0 -1
  289. {zenml_nightly-0.64.0.dev20240811.dist-info → zenml_nightly-0.66.0.dev20240910.dist-info}/LICENSE +0 -0
  290. {zenml_nightly-0.64.0.dev20240811.dist-info → zenml_nightly-0.66.0.dev20240910.dist-info}/WHEEL +0 -0
  291. {zenml_nightly-0.64.0.dev20240811.dist-info → zenml_nightly-0.66.0.dev20240910.dist-info}/entry_points.txt +0 -0
@@ -1,18 +1,18 @@
1
1
  """Utility functions for running pipelines."""
2
2
 
3
- import hashlib
4
3
  import time
5
- from collections import defaultdict
6
4
  from datetime import datetime
7
- from typing import TYPE_CHECKING, Any, Dict, Optional, Set, Tuple, Union
5
+ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Union
8
6
  from uuid import UUID
9
7
 
8
+ from pydantic import BaseModel
9
+
10
10
  from zenml import constants
11
11
  from zenml.client import Client
12
12
  from zenml.config.pipeline_run_configuration import PipelineRunConfiguration
13
- from zenml.config.source import SourceType
13
+ from zenml.config.source import Source, SourceType
14
14
  from zenml.config.step_configurations import StepConfigurationUpdate
15
- from zenml.enums import ExecutionStatus, ModelStages
15
+ from zenml.enums import ExecutionStatus
16
16
  from zenml.logger import get_logger
17
17
  from zenml.models import (
18
18
  FlavorFilter,
@@ -22,15 +22,12 @@ from zenml.models import (
22
22
  PipelineRunResponse,
23
23
  StackResponse,
24
24
  )
25
- from zenml.new.pipelines.model_utils import NewModelRequest
26
25
  from zenml.orchestrators.utils import get_run_name
27
26
  from zenml.stack import Flavor, Stack
28
- from zenml.utils import cloud_utils, code_utils, notebook_utils
27
+ from zenml.utils import code_utils, notebook_utils, source_utils
29
28
  from zenml.zen_stores.base_zen_store import BaseZenStore
30
29
 
31
30
  if TYPE_CHECKING:
32
- from zenml.model.model import Model
33
-
34
31
  StepConfigurationUpdateOrDict = Union[
35
32
  Dict[str, Any], StepConfigurationUpdate
36
33
  ]
@@ -158,116 +155,6 @@ def deploy_pipeline(
158
155
  constants.SHOULD_PREVENT_PIPELINE_EXECUTION = previous_value
159
156
 
160
157
 
161
- def _update_new_requesters(
162
- requester_name: str,
163
- model: "Model",
164
- new_versions_requested: Dict[Tuple[str, Optional[str]], NewModelRequest],
165
- other_models: Set["Model"],
166
- ) -> None:
167
- key = (
168
- model.name,
169
- str(model.version) if model.version else None,
170
- )
171
- if model.version is None:
172
- version_existed = False
173
- else:
174
- try:
175
- model._get_model_version()
176
- version_existed = key not in new_versions_requested
177
- except KeyError as e:
178
- if model.version in ModelStages.values():
179
- raise KeyError(
180
- f"Unable to get model `{model.name}` using stage "
181
- f"`{model.version}`, please check that the model "
182
- "version in given stage exists before running a pipeline."
183
- ) from e
184
- version_existed = False
185
- if not version_existed:
186
- model.was_created_in_this_run = True
187
- new_versions_requested[key].update_request(
188
- model,
189
- NewModelRequest.Requester(source="step", name=requester_name),
190
- )
191
- else:
192
- other_models.add(model)
193
-
194
-
195
- def prepare_model_versions(
196
- deployment: Union["PipelineDeploymentBase", "PipelineDeploymentResponse"],
197
- ) -> None:
198
- """Create model versions which are missing and validate existing ones that are used in the pipeline run.
199
-
200
- Args:
201
- deployment: The pipeline deployment configuration.
202
- """
203
- new_versions_requested: Dict[
204
- Tuple[str, Optional[str]], NewModelRequest
205
- ] = defaultdict(NewModelRequest)
206
- other_models: Set["Model"] = set()
207
- all_steps_have_own_config = True
208
- for step in deployment.step_configurations.values():
209
- step_model = step.config.model
210
- all_steps_have_own_config = (
211
- all_steps_have_own_config and step.config.model is not None
212
- )
213
- if step_model:
214
- _update_new_requesters(
215
- model=step_model,
216
- requester_name=step.config.name,
217
- new_versions_requested=new_versions_requested,
218
- other_models=other_models,
219
- )
220
- if not all_steps_have_own_config:
221
- pipeline_model = deployment.pipeline_configuration.model
222
- if pipeline_model:
223
- _update_new_requesters(
224
- model=pipeline_model,
225
- requester_name=deployment.pipeline_configuration.name,
226
- new_versions_requested=new_versions_requested,
227
- other_models=other_models,
228
- )
229
- elif deployment.pipeline_configuration.model is not None:
230
- logger.warning(
231
- f"ModelConfig of pipeline `{deployment.pipeline_configuration.name}` is overridden in all "
232
- f"steps. "
233
- )
234
-
235
- _validate_new_version_requests(new_versions_requested)
236
-
237
- for other_model in other_models:
238
- other_model._validate_config_in_runtime()
239
-
240
-
241
- def _validate_new_version_requests(
242
- new_versions_requested: Dict[Tuple[str, Optional[str]], NewModelRequest],
243
- ) -> None:
244
- """Validate the model version that are used in the pipeline run.
245
-
246
- Args:
247
- new_versions_requested: A dict of new model version request objects.
248
-
249
- """
250
- is_cloud_model = True
251
- for key, data in new_versions_requested.items():
252
- model_name, model_version = key
253
- if len(data.requesters) > 1:
254
- logger.warning(
255
- f"New version of model version `{model_name}::{model_version or 'NEW'}` "
256
- f"requested in multiple decorators:\n{data.requesters}\n We recommend "
257
- "that `Model` requesting new version is configured only in one "
258
- "place of the pipeline."
259
- )
260
- model_version_response = data.model._validate_config_in_runtime()
261
- is_cloud_model &= cloud_utils.is_cloud_model_version(
262
- model_version_response
263
- )
264
- if not is_cloud_model:
265
- logger.info(
266
- "Models can be viewed in the dashboard using ZenML Pro. Sign up "
267
- "for a free trial at https://www.zenml.io/pro/"
268
- )
269
-
270
-
271
158
  def wait_for_pipeline_run_to_finish(run_id: UUID) -> "PipelineRunResponse":
272
159
  """Waits until a pipeline run is finished.
273
160
 
@@ -383,9 +270,8 @@ def upload_notebook_cell_code_if_necessary(
383
270
  RuntimeError: If the code for one of the steps that will run out of
384
271
  process cannot be extracted into a python file.
385
272
  """
386
- code_archive = code_utils.CodeArchive(root=None)
387
273
  should_upload = False
388
- sources_that_require_upload = []
274
+ resolved_notebook_sources = source_utils.get_resolved_notebook_sources()
389
275
 
390
276
  for step in deployment.step_configurations.values():
391
277
  source = step.spec.source
@@ -396,7 +282,9 @@ def upload_notebook_cell_code_if_necessary(
396
282
  or step.config.step_operator
397
283
  ):
398
284
  should_upload = True
399
- cell_code = getattr(step.spec.source, "_cell_code", None)
285
+ cell_code = resolved_notebook_sources.get(
286
+ source.import_path, None
287
+ )
400
288
 
401
289
  # Code does not run in-process, which means we need to
402
290
  # extract the step code into a python file
@@ -410,20 +298,53 @@ def upload_notebook_cell_code_if_necessary(
410
298
  "of a notebook."
411
299
  )
412
300
 
413
- notebook_utils.warn_about_notebook_cell_magic_commands(
414
- cell_code=cell_code
415
- )
301
+ if should_upload:
302
+ logger.info("Uploading notebook code...")
416
303
 
417
- code_hash = hashlib.sha1(cell_code.encode()).hexdigest() # nosec
418
- module_name = f"extracted_notebook_code_{code_hash}"
419
- file_name = f"{module_name}.py"
420
- code_archive.add_file(source=cell_code, destination=file_name)
304
+ for _, cell_code in resolved_notebook_sources.items():
305
+ notebook_utils.warn_about_notebook_cell_magic_commands(
306
+ cell_code=cell_code
307
+ )
308
+ module_name = notebook_utils.compute_cell_replacement_module_name(
309
+ cell_code=cell_code
310
+ )
311
+ file_name = f"{module_name}.py"
421
312
 
422
- setattr(step.spec.source, "replacement_module", module_name)
423
- sources_that_require_upload.append(source)
313
+ code_utils.upload_notebook_code(
314
+ artifact_store=stack.artifact_store,
315
+ cell_code=cell_code,
316
+ file_name=file_name,
317
+ )
424
318
 
425
- if should_upload:
426
- logger.info("Archiving notebook code...")
427
- code_path = code_utils.upload_code_if_necessary(code_archive)
428
- for source in sources_that_require_upload:
429
- setattr(source, "code_path", code_path)
319
+ all_deployment_sources = get_all_sources_from_value(deployment)
320
+
321
+ for source in all_deployment_sources:
322
+ if source.type == SourceType.NOTEBOOK:
323
+ setattr(source, "artifact_store_id", stack.artifact_store.id)
324
+
325
+ logger.info("Upload finished.")
326
+
327
+
328
+ def get_all_sources_from_value(value: Any) -> List[Source]:
329
+ """Get all source objects from a value.
330
+
331
+ Args:
332
+ value: The value from which to get all the source objects.
333
+
334
+ Returns:
335
+ List of source objects for the given value.
336
+ """
337
+ sources = []
338
+ if isinstance(value, Source):
339
+ sources.append(value)
340
+ elif isinstance(value, BaseModel):
341
+ for v in value.__dict__.values():
342
+ sources.extend(get_all_sources_from_value(v))
343
+ elif isinstance(value, Dict):
344
+ for v in value.values():
345
+ sources.extend(get_all_sources_from_value(v))
346
+ elif isinstance(value, (List, Set, tuple)):
347
+ for v in value:
348
+ sources.extend(get_all_sources_from_value(v))
349
+
350
+ return sources
@@ -118,7 +118,15 @@ class StepContext(metaclass=SingletonMetaClass):
118
118
  """
119
119
  from zenml.client import Client
120
120
 
121
+ try:
122
+ pipeline_run = Client().get_pipeline_run(pipeline_run.id)
123
+ except KeyError:
124
+ pass
121
125
  self.pipeline_run = pipeline_run
126
+ try:
127
+ step_run = Client().get_run_step(step_run.id)
128
+ except KeyError:
129
+ pass
122
130
  self.step_run = step_run
123
131
  self._step_run_info = step_run_info
124
132
  self._cache_enabled = cache_enabled
@@ -177,19 +185,22 @@ class StepContext(metaclass=SingletonMetaClass):
177
185
  Raises:
178
186
  StepContextError: If the `Model` object is not set in `@step` or `@pipeline`.
179
187
  """
180
- if self.step_run.config.model is not None:
181
- model = self.step_run.config.model
188
+ if (
189
+ self.step_run.config.model is not None
190
+ and self.step_run.model_version is not None
191
+ ):
192
+ model = self.step_run.model_version.to_model_class()
182
193
  elif self.pipeline_run.config.model is not None:
183
- model = self.pipeline_run.config.model
194
+ if self.pipeline_run.model_version:
195
+ model = self.pipeline_run.model_version.to_model_class()
196
+ else:
197
+ model = self.pipeline_run.config.model
184
198
  else:
185
199
  raise StepContextError(
186
200
  f"Unable to get Model in step `{self.step_name}` of pipeline "
187
201
  f"run '{self.pipeline_run.id}': it was not set in `@step` or `@pipeline`."
188
202
  )
189
203
 
190
- # warm-up the model version
191
- model._get_or_create_model_version()
192
-
193
204
  return model
194
205
 
195
206
  # TODO: deprecate me
@@ -68,6 +68,15 @@ class BaseOrchestratorConfig(StackComponentConfig):
68
68
  """
69
69
  return False
70
70
 
71
+ @property
72
+ def is_schedulable(self) -> bool:
73
+ """Whether the orchestrator is schedulable or not.
74
+
75
+ Returns:
76
+ Whether the orchestrator is schedulable or not.
77
+ """
78
+ return False
79
+
71
80
 
72
81
  class BaseOrchestrator(StackComponent, ABC):
73
82
  """Base class for all orchestrators.
@@ -225,6 +225,30 @@ class StepLauncher:
225
225
  step_run
226
226
  )
227
227
 
228
+ # warm-up and register model version
229
+ _step_run = None
230
+ model = (
231
+ self._deployment.step_configurations[
232
+ step_run.name
233
+ ].config.model
234
+ or self._deployment.pipeline_configuration.model
235
+ )
236
+ if self._deployment.step_configurations[
237
+ step_run.name
238
+ ].config.model:
239
+ _step_run = step_run_response
240
+
241
+ if model:
242
+ prep_logs_to_show = (
243
+ model._prepare_model_version_before_step_launch(
244
+ pipeline_run=pipeline_run,
245
+ step_run=_step_run,
246
+ return_logs=True,
247
+ )
248
+ )
249
+ if prep_logs_to_show:
250
+ logger.info(prep_logs_to_show)
251
+
228
252
  logger.info(f"Step `{self._step_name}` has started.")
229
253
  if execution_needed:
230
254
  retries = 0
@@ -296,6 +320,17 @@ class StepLauncher:
296
320
  step_run_response.id
297
321
  )
298
322
  raise
323
+ else:
324
+ orchestrator_utils._link_cached_artifacts_to_model(
325
+ model_from_context=model,
326
+ step_run=step_run,
327
+ step_source=self._step.spec.source,
328
+ )
329
+ if model:
330
+ orchestrator_utils._link_pipeline_run_to_model_from_context(
331
+ pipeline_run_id=step_run.pipeline_run_id,
332
+ model=model,
333
+ )
299
334
 
300
335
  except: # noqa: E722
301
336
  logger.error(f"Pipeline run `{pipeline_run.name}` failed.")
@@ -368,10 +403,6 @@ class StepLauncher:
368
403
  Tuple that specifies whether the step needs to be executed as
369
404
  well as the response model of the registered step run.
370
405
  """
371
- model = (
372
- self._deployment.step_configurations[step_run.name].config.model
373
- or self._deployment.pipeline_configuration.model
374
- )
375
406
  input_artifacts, parent_step_ids = input_utils.resolve_step_inputs(
376
407
  step=self._step,
377
408
  run_id=step_run.pipeline_run_id,
@@ -414,20 +445,12 @@ class StepLauncher:
414
445
  execution_needed = False
415
446
  cached_outputs = cached_step_run.outputs
416
447
  step_run.original_step_run_id = cached_step_run.id
448
+
417
449
  step_run.outputs = {
418
450
  output_name: artifact.id
419
451
  for output_name, artifact in cached_outputs.items()
420
452
  }
421
- orchestrator_utils._link_cached_artifacts_to_model(
422
- model_from_context=model,
423
- step_run=step_run,
424
- step_source=self._step.spec.source,
425
- )
426
- if self._step.config.model:
427
- orchestrator_utils._link_pipeline_run_to_model_from_context(
428
- pipeline_run_id=step_run.pipeline_run_id,
429
- model=self._step.config.model,
430
- )
453
+
431
454
  step_run.status = ExecutionStatus.CACHED
432
455
  step_run.end_time = step_run.start_time
433
456
 
@@ -37,7 +37,7 @@ from zenml.constants import (
37
37
  ENV_ZENML_IGNORE_FAILURE_HOOK,
38
38
  handle_bool_env_var,
39
39
  )
40
- from zenml.exceptions import StepContextError, StepInterfaceError
40
+ from zenml.exceptions import StepInterfaceError
41
41
  from zenml.logger import get_logger
42
42
  from zenml.logging.step_logging import StepLogsStorageContext, redirected
43
43
  from zenml.materializers.base_materializer import BaseMaterializer
@@ -181,8 +181,6 @@ class StepRunner:
181
181
  for k, v in output_annotations.items()
182
182
  },
183
183
  )
184
- # Prepare Model Context
185
- self._prepare_model_context_for_step()
186
184
 
187
185
  # Parse the inputs for the entrypoint function.
188
186
  function_params = self._parse_inputs(
@@ -447,14 +445,24 @@ class StepRunner:
447
445
  # we use the datatype of the stored artifact
448
446
  data_type = source_utils.load(artifact.data_type)
449
447
 
448
+ from zenml.orchestrators.utils import (
449
+ register_artifact_store_filesystem,
450
+ )
451
+
450
452
  materializer_class: Type[BaseMaterializer] = (
451
453
  source_utils.load_and_validate_class(
452
454
  artifact.materializer, expected_class=BaseMaterializer
453
455
  )
454
456
  )
455
- materializer: BaseMaterializer = materializer_class(artifact.uri)
456
- materializer.validate_type_compatibility(data_type)
457
- return materializer.load(data_type=data_type)
457
+
458
+ with register_artifact_store_filesystem(
459
+ artifact.artifact_store_id
460
+ ) as target_artifact_store:
461
+ materializer: BaseMaterializer = materializer_class(
462
+ uri=artifact.uri, artifact_store=target_artifact_store
463
+ )
464
+ materializer.validate_type_compatibility(data_type)
465
+ return materializer.load(data_type=data_type)
458
466
 
459
467
  def _validate_outputs(
460
468
  self,
@@ -638,13 +646,6 @@ class StepRunner:
638
646
 
639
647
  return output_artifacts
640
648
 
641
- def _prepare_model_context_for_step(self) -> None:
642
- try:
643
- model = get_step_context().model
644
- model._get_or_create_model_version()
645
- except StepContextError:
646
- return
647
-
648
649
  def load_and_run_hook(
649
650
  self,
650
651
  hook_source: "Source",
@@ -13,9 +13,9 @@
13
13
  # permissions and limitations under the License.
14
14
  """Utility functions for the orchestrator."""
15
15
 
16
+ import os
16
17
  import random
17
- from datetime import datetime
18
- from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple
18
+ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, cast
19
19
  from uuid import UUID
20
20
 
21
21
  from zenml.client import Client
@@ -26,16 +26,21 @@ from zenml.config.source import Source
26
26
  from zenml.constants import (
27
27
  ENV_ZENML_ACTIVE_STACK_ID,
28
28
  ENV_ZENML_ACTIVE_WORKSPACE_ID,
29
+ ENV_ZENML_SERVER,
29
30
  ENV_ZENML_STORE_PREFIX,
30
31
  PIPELINE_API_TOKEN_EXPIRES_MINUTES,
31
32
  )
32
- from zenml.enums import StoreType
33
+ from zenml.enums import StackComponentType, StoreType
33
34
  from zenml.exceptions import StepContextError
35
+ from zenml.logger import get_logger
34
36
  from zenml.model.utils import link_artifact_config_to_model
35
37
  from zenml.models.v2.core.step_run import StepRunRequest
36
38
  from zenml.new.steps.step_context import get_step_context
39
+ from zenml.stack import StackComponent
40
+ from zenml.utils.string_utils import format_name_template
37
41
 
38
42
  if TYPE_CHECKING:
43
+ from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
39
44
  from zenml.artifacts.external_artifact_config import (
40
45
  ExternalArtifactConfiguration,
41
46
  )
@@ -148,10 +153,7 @@ def get_run_name(run_name_template: str) -> str:
148
153
  Returns:
149
154
  The run name derived from the template.
150
155
  """
151
- date = datetime.utcnow().strftime("%Y_%m_%d")
152
- time = datetime.utcnow().strftime("%H_%M_%S_%f")
153
-
154
- run_name = run_name_template.format(date=date, time=time)
156
+ run_name = format_name_template(run_name_template)
155
157
 
156
158
  if run_name == "":
157
159
  raise ValueError("Empty run names are not allowed.")
@@ -305,3 +307,101 @@ def _get_model_versions_from_artifacts(
305
307
  else:
306
308
  break
307
309
  return models
310
+
311
+
312
+ class register_artifact_store_filesystem:
313
+ """Context manager for the artifact_store/filesystem_registry dependency.
314
+
315
+ Even though it is rare, sometimes we bump into cases where we are trying to
316
+ load artifacts that belong to an artifact store which is different from
317
+ the active artifact store.
318
+
319
+ In cases like this, we will try to instantiate the target artifact store
320
+ by creating the corresponding artifact store Python object, which ends up
321
+ registering the right filesystem in the filesystem registry.
322
+
323
+ The problem is, the keys in the filesystem registry are schemes (such as
324
+ "s3://" or "gcs://"). If we have two artifact stores with the same set of
325
+ supported schemes, we might end up overwriting the filesystem that belongs
326
+ to the active artifact store (and its authentication). That's why we have
327
+ to re-instantiate the active artifact store again, so the correct filesystem
328
+ will be restored.
329
+ """
330
+
331
+ def __init__(self, target_artifact_store_id: Optional[UUID]) -> None:
332
+ """Initialization of the context manager.
333
+
334
+ Args:
335
+ target_artifact_store_id: the ID of the artifact store to load.
336
+ """
337
+ self.target_artifact_store_id = target_artifact_store_id
338
+
339
+ def __enter__(self) -> "BaseArtifactStore":
340
+ """Entering the context manager.
341
+
342
+ It creates an instance of the target artifact store to register the
343
+ correct filesystem in the registry.
344
+
345
+ Returns:
346
+ The target artifact store object.
347
+
348
+ Raises:
349
+ RuntimeError: If the target artifact store can not be fetched or
350
+ initiated due to missing dependencies.
351
+ """
352
+ try:
353
+ if self.target_artifact_store_id is not None:
354
+ if (
355
+ Client().active_stack.artifact_store.id
356
+ != self.target_artifact_store_id
357
+ ):
358
+ get_logger(__name__).debug(
359
+ f"Trying to use the artifact store with ID:"
360
+ f"'{self.target_artifact_store_id}'"
361
+ f"which is currently not the active artifact store."
362
+ )
363
+
364
+ artifact_store_model_response = Client().get_stack_component(
365
+ component_type=StackComponentType.ARTIFACT_STORE,
366
+ name_id_or_prefix=self.target_artifact_store_id,
367
+ )
368
+ return cast(
369
+ "BaseArtifactStore",
370
+ StackComponent.from_model(artifact_store_model_response),
371
+ )
372
+ else:
373
+ return Client().active_stack.artifact_store
374
+
375
+ except KeyError:
376
+ raise RuntimeError(
377
+ "Unable to fetch the artifact store with id: "
378
+ f"'{self.target_artifact_store_id}'. Check whether the "
379
+ "artifact store still exists and you have the right "
380
+ "permissions to access it."
381
+ )
382
+ except ImportError:
383
+ raise RuntimeError(
384
+ "Unable to load the implementation of the artifact store with"
385
+ f"id: '{self.target_artifact_store_id}'. Please make sure that "
386
+ "the environment that you are loading this artifact from "
387
+ "has the right dependencies."
388
+ )
389
+
390
+ def __exit__(
391
+ self,
392
+ exc_type: Optional[Any],
393
+ exc_value: Optional[Any],
394
+ traceback: Optional[Any],
395
+ ) -> None:
396
+ """Set it back to the original state.
397
+
398
+ Args:
399
+ exc_type: The class of the exception
400
+ exc_value: The instance of the exception
401
+ traceback: The traceback of the exception
402
+ """
403
+ if ENV_ZENML_SERVER not in os.environ:
404
+ # As we exit the handler, we have to re-register the filesystem
405
+ # that belongs to the active artifact store as it may have been
406
+ # overwritten.
407
+ Client().active_stack.artifact_store._register()
@@ -18,10 +18,10 @@ from uuid import UUID
18
18
 
19
19
  from zenml.client import Client
20
20
  from zenml.enums import StackComponentType
21
- from zenml.models.v2.core.service_connector import ServiceConnectorRequest
22
- from zenml.models.v2.misc.full_stack import (
21
+ from zenml.models import (
23
22
  ResourcesInfo,
24
23
  ServiceConnectorInfo,
24
+ ServiceConnectorRequest,
25
25
  ServiceConnectorResourcesInfo,
26
26
  )
27
27
  from zenml.utils.pagination_utils import depaginate
zenml/stack/utils.py CHANGED
@@ -13,7 +13,7 @@
13
13
  # permissions and limitations under the License.
14
14
  """Util functions for handling stacks, components, and flavors."""
15
15
 
16
- from typing import Any, Dict, Optional
16
+ from typing import Any, Dict, Optional, Type
17
17
 
18
18
  from zenml.client import Client
19
19
  from zenml.enums import StackComponentType, StoreType
@@ -75,7 +75,16 @@ def validate_stack_component_config(
75
75
  return None
76
76
  raise
77
77
 
78
- configuration = flavor_class.config_class(**configuration_dict)
78
+ config_class = flavor_class.config_class
79
+ # Make sure extras are forbidden for the config class. Due to inheritance
80
+ # order, some config classes allow extras by accident which we patch here.
81
+ validation_config_class: Type[StackComponentConfig] = type(
82
+ config_class.__name__,
83
+ (config_class,),
84
+ {"model_config": {"extra": "forbid"}},
85
+ )
86
+ configuration = validation_config_class(**configuration_dict)
87
+
79
88
  if not configuration.is_valid:
80
89
  raise ValueError(
81
90
  f"Invalid stack component configuration. Please verify "
@@ -69,7 +69,8 @@ permissions and are aware of any potential costs:
69
69
  - An Azure Resource Group to contain all the resources required for the ZenML stack
70
70
  - An Azure Storage Account and Blob Storage Container registered as a [ZenML artifact store](https://docs.zenml.io/stack-components/artifact-stores/azure).
71
71
  - An Azure Container Registry registered as a [ZenML container registry](https://docs.zenml.io/stack-components/container-registries/azure).
72
- - An AzureML Workspace registered as a [ZenML orchestrator](https://docs.zenml.io/stack-components/orchestrators/azureml) and used to run pipelines.
72
+ - An AzureML Workspace registered as both a [ZenML orchestrator](https://docs.zenml.io/stack-components/orchestrators/azureml) and a
73
+ [ZenML step operator](https://docs.zenml.io/stack-components/step-operators/azureml) and used to run pipelines.
73
74
  A Key Vault and Application Insights instance will also be created in the same Resource Group and used to construct the AzureML Workspace.
74
75
  - An Azure Service Principal with the minimum necessary permissions to access
75
76
  the above resources.