zenml-nightly 0.58.2.dev20240626__py3-none-any.whl → 0.62.0.dev20240726__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 (266) hide show
  1. README.md +31 -10
  2. RELEASE_NOTES.md +280 -0
  3. zenml/VERSION +1 -1
  4. zenml/__init__.py +2 -0
  5. zenml/analytics/enums.py +3 -0
  6. zenml/cli/__init__.py +28 -0
  7. zenml/cli/artifact.py +1 -2
  8. zenml/cli/integration.py +9 -8
  9. zenml/cli/server.py +6 -0
  10. zenml/cli/stack.py +812 -39
  11. zenml/cli/stack_components.py +9 -0
  12. zenml/cli/text_utils.py +35 -1
  13. zenml/cli/utils.py +127 -10
  14. zenml/client.py +23 -14
  15. zenml/config/docker_settings.py +8 -5
  16. zenml/constants.py +13 -1
  17. zenml/container_registries/base_container_registry.py +1 -0
  18. zenml/enums.py +23 -0
  19. zenml/event_hub/event_hub.py +5 -8
  20. zenml/integrations/__init__.py +1 -0
  21. zenml/integrations/aws/__init__.py +1 -0
  22. zenml/integrations/azure/__init__.py +3 -2
  23. zenml/integrations/constants.py +1 -0
  24. zenml/integrations/databricks/__init__.py +52 -0
  25. zenml/integrations/databricks/flavors/__init__.py +30 -0
  26. zenml/integrations/databricks/flavors/databricks_model_deployer_flavor.py +118 -0
  27. zenml/integrations/databricks/flavors/databricks_orchestrator_flavor.py +147 -0
  28. zenml/integrations/databricks/model_deployers/__init__.py +20 -0
  29. zenml/integrations/databricks/model_deployers/databricks_model_deployer.py +249 -0
  30. zenml/integrations/databricks/orchestrators/__init__.py +20 -0
  31. zenml/integrations/databricks/orchestrators/databricks_orchestrator.py +497 -0
  32. zenml/integrations/databricks/orchestrators/databricks_orchestrator_entrypoint_config.py +97 -0
  33. zenml/integrations/databricks/services/__init__.py +19 -0
  34. zenml/integrations/databricks/services/databricks_deployment.py +407 -0
  35. zenml/integrations/databricks/utils/__init__.py +14 -0
  36. zenml/integrations/databricks/utils/databricks_utils.py +87 -0
  37. zenml/integrations/deepchecks/__init__.py +1 -0
  38. zenml/integrations/discord/__init__.py +1 -0
  39. zenml/integrations/evidently/__init__.py +1 -0
  40. zenml/integrations/facets/__init__.py +1 -0
  41. zenml/integrations/feast/__init__.py +1 -0
  42. zenml/integrations/gcp/__init__.py +3 -1
  43. zenml/integrations/gcp/google_credentials_mixin.py +1 -1
  44. zenml/integrations/gcp/service_connectors/gcp_service_connector.py +320 -64
  45. zenml/integrations/great_expectations/data_validators/ge_data_validator.py +12 -8
  46. zenml/integrations/huggingface/__init__.py +1 -0
  47. zenml/integrations/huggingface/materializers/huggingface_datasets_materializer.py +88 -3
  48. zenml/integrations/huggingface/steps/accelerate_runner.py +1 -7
  49. zenml/integrations/integration.py +24 -0
  50. zenml/integrations/kubeflow/__init__.py +3 -0
  51. zenml/integrations/kubeflow/flavors/kubeflow_orchestrator_flavor.py +1 -1
  52. zenml/integrations/kubeflow/orchestrators/kubeflow_orchestrator.py +0 -1
  53. zenml/integrations/kubernetes/__init__.py +3 -1
  54. zenml/integrations/kubernetes/orchestrators/kube_utils.py +4 -1
  55. zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py +1 -13
  56. zenml/integrations/kubernetes/orchestrators/manifest_utils.py +22 -4
  57. zenml/integrations/kubernetes/pod_settings.py +4 -0
  58. zenml/integrations/label_studio/annotators/label_studio_annotator.py +1 -0
  59. zenml/integrations/langchain/__init__.py +1 -0
  60. zenml/integrations/lightgbm/__init__.py +1 -0
  61. zenml/integrations/mlflow/__init__.py +4 -2
  62. zenml/integrations/mlflow/model_registries/mlflow_model_registry.py +6 -2
  63. zenml/integrations/mlflow/services/mlflow_deployment.py +1 -1
  64. zenml/integrations/neural_prophet/__init__.py +1 -0
  65. zenml/integrations/polars/__init__.py +1 -0
  66. zenml/integrations/prodigy/__init__.py +1 -0
  67. zenml/integrations/pycaret/__init__.py +6 -0
  68. zenml/integrations/registry.py +37 -0
  69. zenml/integrations/s3/artifact_stores/s3_artifact_store.py +17 -6
  70. zenml/integrations/seldon/__init__.py +1 -0
  71. zenml/integrations/seldon/model_deployers/seldon_model_deployer.py +1 -0
  72. zenml/integrations/skypilot/flavors/skypilot_orchestrator_base_vm_config.py +2 -2
  73. zenml/integrations/skypilot/orchestrators/skypilot_base_vm_orchestrator.py +1 -1
  74. zenml/integrations/skypilot/orchestrators/skypilot_orchestrator_entrypoint.py +2 -2
  75. zenml/integrations/skypilot_aws/__init__.py +2 -1
  76. zenml/integrations/skypilot_azure/__init__.py +1 -3
  77. zenml/integrations/skypilot_gcp/__init__.py +1 -1
  78. zenml/integrations/skypilot_lambda/__init__.py +1 -1
  79. zenml/integrations/skypilot_lambda/flavors/skypilot_orchestrator_lambda_vm_flavor.py +1 -1
  80. zenml/integrations/slack/__init__.py +1 -0
  81. zenml/integrations/tekton/__init__.py +1 -0
  82. zenml/integrations/tensorboard/__init__.py +0 -1
  83. zenml/integrations/tensorflow/__init__.py +18 -6
  84. zenml/integrations/wandb/__init__.py +1 -0
  85. zenml/logging/step_logging.py +34 -35
  86. zenml/materializers/built_in_materializer.py +1 -1
  87. zenml/materializers/cloudpickle_materializer.py +1 -1
  88. zenml/model/model.py +1 -1
  89. zenml/models/__init__.py +11 -0
  90. zenml/models/v2/core/component.py +47 -0
  91. zenml/models/v2/core/model.py +1 -2
  92. zenml/models/v2/core/server_settings.py +0 -20
  93. zenml/models/v2/core/service_connector.py +17 -0
  94. zenml/models/v2/core/stack.py +31 -0
  95. zenml/models/v2/misc/full_stack.py +129 -0
  96. zenml/models/v2/misc/stack_deployment.py +91 -0
  97. zenml/new/pipelines/pipeline.py +1 -1
  98. zenml/new/pipelines/run_utils.py +1 -1
  99. zenml/orchestrators/__init__.py +4 -0
  100. zenml/orchestrators/input_utils.py +3 -6
  101. zenml/orchestrators/step_launcher.py +1 -0
  102. zenml/orchestrators/wheeled_orchestrator.py +147 -0
  103. zenml/service_connectors/service_connector_utils.py +408 -0
  104. zenml/stack/stack.py +3 -6
  105. zenml/stack_deployments/__init__.py +14 -0
  106. zenml/stack_deployments/aws_stack_deployment.py +254 -0
  107. zenml/stack_deployments/azure_stack_deployment.py +179 -0
  108. zenml/stack_deployments/gcp_stack_deployment.py +269 -0
  109. zenml/stack_deployments/stack_deployment.py +218 -0
  110. zenml/stack_deployments/utils.py +48 -0
  111. zenml/steps/base_step.py +7 -5
  112. zenml/utils/function_utils.py +2 -2
  113. zenml/utils/pagination_utils.py +7 -5
  114. zenml/utils/pipeline_docker_image_builder.py +105 -68
  115. zenml/utils/pydantic_utils.py +6 -5
  116. zenml/utils/source_utils.py +4 -1
  117. zenml/zen_server/cloud_utils.py +18 -3
  118. zenml/zen_server/dashboard/assets/{404-CDPQCl4D.js → 404-B_YdvmwS.js} +1 -1
  119. zenml/zen_server/dashboard/assets/@radix-CFOkMR_E.js +85 -0
  120. zenml/zen_server/dashboard/assets/{@react-router-DYovave8.js → @react-router-CO-OsFwI.js} +2 -2
  121. zenml/zen_server/dashboard/assets/{@reactflow-CHBapDaj.js → @reactflow-l_1hUr1S.js} +2 -2
  122. zenml/zen_server/dashboard/assets/@tanstack-DYiOyJUL.js +22 -0
  123. zenml/zen_server/dashboard/assets/AwarenessChannel-CFg5iX4Z.js +1 -0
  124. zenml/zen_server/dashboard/assets/{CodeSnippet-BidtnWOi.js → CodeSnippet-Dvkx_82E.js} +2 -2
  125. zenml/zen_server/dashboard/assets/CollapsibleCard-opiuBHHc.js +1 -0
  126. zenml/zen_server/dashboard/assets/Commands-DoN1xrEq.js +1 -0
  127. zenml/zen_server/dashboard/assets/CopyButton-Cr7xYEPb.js +2 -0
  128. zenml/zen_server/dashboard/assets/{CsvVizualization-BOuez-fG.js → CsvVizualization-Ck-nZ43m.js} +7 -7
  129. zenml/zen_server/dashboard/assets/DisplayDate-DYgIjlDF.js +1 -0
  130. zenml/zen_server/dashboard/assets/EmptyState-BMLnFVlB.js +1 -0
  131. zenml/zen_server/dashboard/assets/Error-kLtljEOM.js +1 -0
  132. zenml/zen_server/dashboard/assets/ExecutionStatus-DguLLgTK.js +1 -0
  133. zenml/zen_server/dashboard/assets/Helpbox-BXUMP21n.js +1 -0
  134. zenml/zen_server/dashboard/assets/Infobox-DSt0O-dm.js +1 -0
  135. zenml/zen_server/dashboard/assets/InlineAvatar-xsrsIGE-.js +1 -0
  136. zenml/zen_server/dashboard/assets/{MarkdownVisualization-DsB2QZiK.js → MarkdownVisualization-xp3hhULl.js} +2 -2
  137. zenml/zen_server/dashboard/assets/Pagination-C6X-mifw.js +1 -0
  138. zenml/zen_server/dashboard/assets/PasswordChecker-DUveqlva.js +1 -0
  139. zenml/zen_server/dashboard/assets/SetPassword-BXGTWiwj.js +1 -0
  140. zenml/zen_server/dashboard/assets/SuccessStep-DZC60t0x.js +1 -0
  141. zenml/zen_server/dashboard/assets/{UpdatePasswordSchemas-DnM-c11H.js → UpdatePasswordSchemas-DGvwFWO1.js} +1 -1
  142. zenml/zen_server/dashboard/assets/{aws-t0gKCj_R.js → aws-BgKTfTfx.js} +1 -1
  143. zenml/zen_server/dashboard/assets/{check-circle-BVvhm5dy.js → check-circle-i56092KI.js} +1 -1
  144. zenml/zen_server/dashboard/assets/{chevron-right-double-CJ50E9Gr.js → chevron-right-double-CZBOf6JM.js} +1 -1
  145. zenml/zen_server/dashboard/assets/cloud-only-C_yFCAkP.js +1 -0
  146. zenml/zen_server/dashboard/assets/{copy-BRhQz3j-.js → copy-BXNk6BjL.js} +1 -1
  147. zenml/zen_server/dashboard/assets/{database-CRRnyFWh.js → database-1xWSgZfO.js} +1 -1
  148. zenml/zen_server/dashboard/assets/{docker-BAonhm6G.js → docker-CQMVm_4d.js} +1 -1
  149. zenml/zen_server/dashboard/assets/{file-text-CbVERUON.js → file-text-CqD_iu6l.js} +1 -1
  150. zenml/zen_server/dashboard/assets/{help-B8rqCvqn.js → help-bu_DgLKI.js} +1 -1
  151. zenml/zen_server/dashboard/assets/index-BczVOqUf.js +55 -0
  152. zenml/zen_server/dashboard/assets/index-EpMIKgrI.css +1 -0
  153. zenml/zen_server/dashboard/assets/index-rK_Wuy2W.js +1 -0
  154. zenml/zen_server/dashboard/assets/index.esm-Corw4lXQ.js +1 -0
  155. zenml/zen_server/dashboard/assets/{login-mutation-wzzl23C6.js → login-mutation-CrHrndTI.js} +1 -1
  156. zenml/zen_server/dashboard/assets/logs-D8k8BVFf.js +1 -0
  157. zenml/zen_server/dashboard/assets/not-found-DYa4pC-C.js +1 -0
  158. zenml/zen_server/dashboard/assets/package-B3fWP-Dh.js +1 -0
  159. zenml/zen_server/dashboard/assets/page-1h_sD1jz.js +1 -0
  160. zenml/zen_server/dashboard/assets/{page-yN4rZ-ZS.js → page-1iL8aMqs.js} +1 -1
  161. zenml/zen_server/dashboard/assets/{page-Bi5AI0S7.js → page-2grKx_MY.js} +1 -1
  162. zenml/zen_server/dashboard/assets/page-5NCOHOsy.js +1 -0
  163. zenml/zen_server/dashboard/assets/page-8a4UMKXZ.js +1 -0
  164. zenml/zen_server/dashboard/assets/{page-AQKopn_4.js → page-B6h3iaHJ.js} +1 -1
  165. zenml/zen_server/dashboard/assets/page-BDns21Iz.js +1 -0
  166. zenml/zen_server/dashboard/assets/{page-BmkSiYeQ.js → page-BhgCDInH.js} +2 -2
  167. zenml/zen_server/dashboard/assets/{page-BzVZGExK.js → page-Bi-wtWiO.js} +2 -2
  168. zenml/zen_server/dashboard/assets/page-BkeAAYwp.js +1 -0
  169. zenml/zen_server/dashboard/assets/page-BkuQDIf-.js +1 -0
  170. zenml/zen_server/dashboard/assets/page-BnaevhnB.js +1 -0
  171. zenml/zen_server/dashboard/assets/page-Bq0YxkLV.js +1 -0
  172. zenml/zen_server/dashboard/assets/page-Bs2F4eoD.js +2 -0
  173. zenml/zen_server/dashboard/assets/page-C6-UGEbH.js +1 -0
  174. zenml/zen_server/dashboard/assets/page-CCNRIt_f.js +1 -0
  175. zenml/zen_server/dashboard/assets/page-CHNxpz3n.js +1 -0
  176. zenml/zen_server/dashboard/assets/page-DgorQFqi.js +1 -0
  177. zenml/zen_server/dashboard/assets/page-K8ebxVIs.js +1 -0
  178. zenml/zen_server/dashboard/assets/{page-CuT1SUik.js → page-MFQyIJd3.js} +1 -1
  179. zenml/zen_server/dashboard/assets/page-TgCF0P_U.js +1 -0
  180. zenml/zen_server/dashboard/assets/page-ZnCEe-eK.js +9 -0
  181. zenml/zen_server/dashboard/assets/{page-BW6Ket3a.js → page-uA5prJGY.js} +1 -1
  182. zenml/zen_server/dashboard/assets/persist-D7HJNBWx.js +1 -0
  183. zenml/zen_server/dashboard/assets/{play-circle-DK5QMJyp.js → play-circle-CNtZKDnW.js} +1 -1
  184. zenml/zen_server/dashboard/assets/plus-C8WOyCzt.js +1 -0
  185. zenml/zen_server/dashboard/assets/stack-detail-query-Cficsl6d.js +1 -0
  186. zenml/zen_server/dashboard/assets/{terminal-B2ovgWuz.js → terminal-By9cErXc.js} +1 -1
  187. zenml/zen_server/dashboard/assets/update-server-settings-mutation-7d8xi1tS.js +1 -0
  188. zenml/zen_server/dashboard/assets/{url-6_xv0WJS.js → url-D7mAQGUM.js} +1 -1
  189. zenml/zen_server/dashboard/assets/{zod-DrZvVLjd.js → zod-BhoGpZ63.js} +1 -1
  190. zenml/zen_server/dashboard/index.html +7 -7
  191. zenml/zen_server/dashboard_legacy/asset-manifest.json +4 -4
  192. zenml/zen_server/dashboard_legacy/index.html +1 -1
  193. zenml/zen_server/dashboard_legacy/{precache-manifest.f4abc5b7cfa7d90c1caf5521918e29a8.js → precache-manifest.12246c7548e71e2c4438e496360de80c.js} +4 -4
  194. zenml/zen_server/dashboard_legacy/service-worker.js +1 -1
  195. zenml/zen_server/dashboard_legacy/static/js/main.3b27024b.chunk.js +2 -0
  196. zenml/zen_server/dashboard_legacy/static/js/{main.ac2f17d0.chunk.js.map → main.3b27024b.chunk.js.map} +1 -1
  197. zenml/zen_server/deploy/helm/Chart.yaml +1 -1
  198. zenml/zen_server/deploy/helm/README.md +2 -2
  199. zenml/zen_server/feature_gate/zenml_cloud_feature_gate.py +11 -5
  200. zenml/zen_server/pipeline_deployment/utils.py +57 -44
  201. zenml/zen_server/rbac/utils.py +10 -2
  202. zenml/zen_server/rbac/zenml_cloud_rbac.py +11 -5
  203. zenml/zen_server/routers/devices_endpoints.py +4 -1
  204. zenml/zen_server/routers/server_endpoints.py +29 -2
  205. zenml/zen_server/routers/service_connectors_endpoints.py +57 -0
  206. zenml/zen_server/routers/stack_deployment_endpoints.py +158 -0
  207. zenml/zen_server/routers/steps_endpoints.py +2 -1
  208. zenml/zen_server/routers/workspaces_endpoints.py +64 -0
  209. zenml/zen_server/zen_server_api.py +2 -0
  210. zenml/zen_stores/migrations/utils.py +1 -1
  211. zenml/zen_stores/migrations/versions/0.60.0_release.py +23 -0
  212. zenml/zen_stores/migrations/versions/0.61.0_release.py +23 -0
  213. zenml/zen_stores/migrations/versions/0.62.0_release.py +23 -0
  214. zenml/zen_stores/migrations/versions/0d707865f404_adding_labels_to_stacks.py +30 -0
  215. zenml/zen_stores/migrations/versions/b4fca5241eea_migrate_onboarding_state.py +167 -0
  216. zenml/zen_stores/rest_zen_store.py +149 -4
  217. zenml/zen_stores/schemas/component_schemas.py +14 -0
  218. zenml/zen_stores/schemas/server_settings_schemas.py +23 -11
  219. zenml/zen_stores/schemas/stack_schemas.py +10 -0
  220. zenml/zen_stores/schemas/step_run_schemas.py +27 -11
  221. zenml/zen_stores/sql_zen_store.py +450 -6
  222. zenml/zen_stores/zen_store_interface.py +80 -0
  223. {zenml_nightly-0.58.2.dev20240626.dist-info → zenml_nightly-0.62.0.dev20240726.dist-info}/METADATA +35 -13
  224. {zenml_nightly-0.58.2.dev20240626.dist-info → zenml_nightly-0.62.0.dev20240726.dist-info}/RECORD +227 -191
  225. zenml/zen_server/dashboard/assets/@radix-C9DBgJhe.js +0 -77
  226. zenml/zen_server/dashboard/assets/@tanstack-CEbkxrhX.js +0 -30
  227. zenml/zen_server/dashboard/assets/AwarenessChannel-nXGpmj_f.js +0 -1
  228. zenml/zen_server/dashboard/assets/Cards-nwsvQLVS.js +0 -1
  229. zenml/zen_server/dashboard/assets/Commands-DuIWKg_Q.js +0 -1
  230. zenml/zen_server/dashboard/assets/CopyButton-B_YSm-Ds.js +0 -2
  231. zenml/zen_server/dashboard/assets/DisplayDate-BdguISQF.js +0 -1
  232. zenml/zen_server/dashboard/assets/EmptyState-BkooiGtL.js +0 -1
  233. zenml/zen_server/dashboard/assets/Error-B6M0dPph.js +0 -1
  234. zenml/zen_server/dashboard/assets/Helpbox-BQoqCm04.js +0 -1
  235. zenml/zen_server/dashboard/assets/Infobox-Ce9mefqU.js +0 -1
  236. zenml/zen_server/dashboard/assets/InlineAvatar-DGf3dVhV.js +0 -1
  237. zenml/zen_server/dashboard/assets/PageHeader-DGaemzjc.js +0 -1
  238. zenml/zen_server/dashboard/assets/Pagination-DVYfBCCc.js +0 -1
  239. zenml/zen_server/dashboard/assets/PasswordChecker-DSLBp7Vl.js +0 -1
  240. zenml/zen_server/dashboard/assets/SetPassword-B5s7DJug.js +0 -1
  241. zenml/zen_server/dashboard/assets/SuccessStep-ZzczaM7g.js +0 -1
  242. zenml/zen_server/dashboard/assets/chevron-down-zcvCWmyP.js +0 -1
  243. zenml/zen_server/dashboard/assets/cloud-only-Ba_ShBR5.js +0 -1
  244. zenml/zen_server/dashboard/assets/index-CWJ3xbIf.css +0 -1
  245. zenml/zen_server/dashboard/assets/index-QORVVTMN.js +0 -55
  246. zenml/zen_server/dashboard/assets/index.esm-F7nqy9zY.js +0 -1
  247. zenml/zen_server/dashboard/assets/not-found-Dh2la7kh.js +0 -1
  248. zenml/zen_server/dashboard/assets/page-B-5jAKoO.js +0 -1
  249. zenml/zen_server/dashboard/assets/page-B-vWk8a6.js +0 -1
  250. zenml/zen_server/dashboard/assets/page-B0BrqfS8.js +0 -1
  251. zenml/zen_server/dashboard/assets/page-BQxVFlUl.js +0 -1
  252. zenml/zen_server/dashboard/assets/page-ByrHy6Ss.js +0 -1
  253. zenml/zen_server/dashboard/assets/page-CPtY4Kv_.js +0 -1
  254. zenml/zen_server/dashboard/assets/page-CmmukLsl.js +0 -1
  255. zenml/zen_server/dashboard/assets/page-D2D-7qyr.js +0 -9
  256. zenml/zen_server/dashboard/assets/page-DAQQyLxT.js +0 -1
  257. zenml/zen_server/dashboard/assets/page-DHkUMl_E.js +0 -1
  258. zenml/zen_server/dashboard/assets/page-DZCbwOEs.js +0 -2
  259. zenml/zen_server/dashboard/assets/page-DdaIt20-.js +0 -1
  260. zenml/zen_server/dashboard/assets/page-LqLs24Ot.js +0 -1
  261. zenml/zen_server/dashboard/assets/page-lebv0c7C.js +0 -1
  262. zenml/zen_server/dashboard/assets/update-server-settings-mutation-0Wgz8pUE.js +0 -1
  263. zenml/zen_server/dashboard_legacy/static/js/main.ac2f17d0.chunk.js +0 -2
  264. {zenml_nightly-0.58.2.dev20240626.dist-info → zenml_nightly-0.62.0.dev20240726.dist-info}/LICENSE +0 -0
  265. {zenml_nightly-0.58.2.dev20240626.dist-info → zenml_nightly-0.62.0.dev20240726.dist-info}/WHEEL +0 -0
  266. {zenml_nightly-0.58.2.dev20240626.dist-info → zenml_nightly-0.62.0.dev20240726.dist-info}/entry_points.txt +0 -0
@@ -63,6 +63,7 @@ from zenml.zen_server.routers import (
63
63
  service_connectors_endpoints,
64
64
  service_endpoints,
65
65
  stack_components_endpoints,
66
+ stack_deployment_endpoints,
66
67
  stacks_endpoints,
67
68
  steps_endpoints,
68
69
  tags_endpoints,
@@ -290,6 +291,7 @@ app.include_router(service_accounts_endpoints.router)
290
291
  app.include_router(service_connectors_endpoints.router)
291
292
  app.include_router(service_connectors_endpoints.types_router)
292
293
  app.include_router(service_endpoints.router)
294
+ app.include_router(stack_deployment_endpoints.router)
293
295
  app.include_router(stacks_endpoints.router)
294
296
  app.include_router(stack_components_endpoints.router)
295
297
  app.include_router(stack_components_endpoints.types_router)
@@ -364,7 +364,7 @@ class MigrationUtils(BaseModel):
364
364
  # Convert column values to the correct type
365
365
  for column in table.columns:
366
366
  # Blob columns are stored as binary strings
367
- if column.type.python_type == bytes and isinstance(
367
+ if column.type.python_type is bytes and isinstance(
368
368
  row[column.name], str
369
369
  ):
370
370
  # Convert the string to bytes
@@ -0,0 +1,23 @@
1
+ """Release [0.60.0].
2
+
3
+ Revision ID: 0.60.0
4
+ Revises: 25155145c545
5
+ Create Date: 2024-06-25 14:26:48.867005
6
+
7
+ """
8
+
9
+ # revision identifiers, used by Alembic.
10
+ revision = "0.60.0"
11
+ down_revision = "25155145c545"
12
+ branch_labels = None
13
+ depends_on = None
14
+
15
+
16
+ def upgrade() -> None:
17
+ """Upgrade database schema and/or data, creating a new revision."""
18
+ pass
19
+
20
+
21
+ def downgrade() -> None:
22
+ """Downgrade database schema and/or data back to the previous revision."""
23
+ pass
@@ -0,0 +1,23 @@
1
+ """Release [0.61.0].
2
+
3
+ Revision ID: 0.61.0
4
+ Revises: 0d707865f404
5
+ Create Date: 2024-07-08 15:52:54.765307
6
+
7
+ """
8
+
9
+ # revision identifiers, used by Alembic.
10
+ revision = "0.61.0"
11
+ down_revision = "0d707865f404"
12
+ branch_labels = None
13
+ depends_on = None
14
+
15
+
16
+ def upgrade() -> None:
17
+ """Upgrade database schema and/or data, creating a new revision."""
18
+ pass
19
+
20
+
21
+ def downgrade() -> None:
22
+ """Downgrade database schema and/or data back to the previous revision."""
23
+ pass
@@ -0,0 +1,23 @@
1
+ """Release [0.62.0].
2
+
3
+ Revision ID: 0.62.0
4
+ Revises: b4fca5241eea
5
+ Create Date: 2024-07-15 14:15:45.347033
6
+
7
+ """
8
+
9
+ # revision identifiers, used by Alembic.
10
+ revision = "0.62.0"
11
+ down_revision = "b4fca5241eea"
12
+ branch_labels = None
13
+ depends_on = None
14
+
15
+
16
+ def upgrade() -> None:
17
+ """Upgrade database schema and/or data, creating a new revision."""
18
+ pass
19
+
20
+
21
+ def downgrade() -> None:
22
+ """Downgrade database schema and/or data back to the previous revision."""
23
+ pass
@@ -0,0 +1,30 @@
1
+ """adding labels to stacks [0d707865f404].
2
+
3
+ Revision ID: 0d707865f404
4
+ Revises: 0.60.0
5
+ Create Date: 2024-07-04 16:10:07.709184
6
+
7
+ """
8
+
9
+ import sqlalchemy as sa
10
+ from alembic import op
11
+
12
+ # revision identifiers, used by Alembic.
13
+ revision = "0d707865f404"
14
+ down_revision = "0.60.0"
15
+ branch_labels = None
16
+ depends_on = None
17
+
18
+
19
+ def upgrade() -> None:
20
+ """Upgrade database schema and/or data, creating a new revision."""
21
+ with op.batch_alter_table("stack", schema=None) as batch_op:
22
+ batch_op.add_column(
23
+ sa.Column("labels", sa.LargeBinary(), nullable=True)
24
+ )
25
+
26
+
27
+ def downgrade() -> None:
28
+ """Downgrade database schema and/or data back to the previous revision."""
29
+ with op.batch_alter_table("stack", schema=None) as batch_op:
30
+ batch_op.drop_column("labels")
@@ -0,0 +1,167 @@
1
+ """Migrate onboarding state [b4fca5241eea].
2
+
3
+ Revision ID: b4fca5241eea
4
+ Revises: 0.61.0
5
+ Create Date: 2024-06-20 15:01:22.414801
6
+
7
+ """
8
+
9
+ import json
10
+ from typing import Dict, List
11
+
12
+ import sqlalchemy as sa
13
+ from alembic import op
14
+
15
+ # revision identifiers, used by Alembic.
16
+ revision = "b4fca5241eea"
17
+ down_revision = "0.61.0"
18
+ branch_labels = None
19
+ depends_on = None
20
+
21
+
22
+ ONBOARDING_KEY_MAPPING = {
23
+ "connect_zenml": ["device_verified"],
24
+ "run_first_pipeline": ["pipeline_run", "starter_setup_completed"],
25
+ "run_remote_pipeline": [
26
+ "production_setup_completed",
27
+ ],
28
+ }
29
+
30
+
31
+ def upgrade() -> None:
32
+ """Upgrade database schema and/or data, creating a new revision."""
33
+ connection = op.get_bind()
34
+
35
+ meta = sa.MetaData()
36
+ meta.reflect(only=("server_settings",), bind=connection)
37
+
38
+ server_settings_table = sa.Table("server_settings", meta)
39
+
40
+ existing_onboarding_state = connection.execute(
41
+ sa.select(server_settings_table.c.onboarding_state)
42
+ ).scalar_one_or_none()
43
+
44
+ new_state = []
45
+
46
+ if existing_onboarding_state:
47
+ # There was already an existing onboarding state in the DB
48
+ # -> Migrate to the new server keys
49
+ state = json.loads(existing_onboarding_state)
50
+
51
+ if isinstance(state, Dict):
52
+ for key in state.keys():
53
+ if key in ONBOARDING_KEY_MAPPING:
54
+ new_state.extend(ONBOARDING_KEY_MAPPING[key])
55
+ elif isinstance(state, List):
56
+ # Somehow the state is already converted, probably shouldn't happen
57
+ return
58
+
59
+ # We now query the DB and complete all onboarding steps that we can detect
60
+ # from the database
61
+ meta = sa.MetaData()
62
+ meta.reflect(
63
+ only=(
64
+ "pipeline_run",
65
+ "stack_component",
66
+ "stack",
67
+ "stack_composition",
68
+ "pipeline_deployment",
69
+ ),
70
+ bind=connection,
71
+ )
72
+
73
+ pipeline_run_table = sa.Table("pipeline_run", meta)
74
+ stack_component_table = sa.Table("stack_component", meta)
75
+ stack_table = sa.Table("stack", meta)
76
+ stack_composition_table = sa.Table("stack_composition", meta)
77
+ pipeline_deployment_table = sa.Table("pipeline_deployment", meta)
78
+
79
+ pipeline_run_count = connection.execute(
80
+ sa.select(sa.func.count(pipeline_run_table.c.id))
81
+ ).scalar()
82
+ if pipeline_run_count and pipeline_run_count > 0:
83
+ new_state.extend(ONBOARDING_KEY_MAPPING["run_first_pipeline"])
84
+
85
+ stack_with_remote_orchestrator_count = connection.execute(
86
+ sa.select(sa.func.count(stack_table.c.id))
87
+ .where(stack_composition_table.c.stack_id == stack_table.c.id)
88
+ .where(
89
+ stack_composition_table.c.component_id
90
+ == stack_component_table.c.id
91
+ )
92
+ .where(
93
+ stack_component_table.c.flavor.not_in(["local", "local_docker"])
94
+ )
95
+ .where(stack_component_table.c.type == "orchestrator")
96
+ ).scalar()
97
+ if (
98
+ stack_with_remote_orchestrator_count
99
+ and stack_with_remote_orchestrator_count > 0
100
+ ):
101
+ new_state.append("stack_with_remote_orchestrator_created")
102
+
103
+ pipeline_run_with_remote_artifact_store_count = connection.execute(
104
+ sa.select(sa.func.count(pipeline_run_table.c.id))
105
+ .where(
106
+ pipeline_run_table.c.deployment_id
107
+ == pipeline_deployment_table.c.id
108
+ )
109
+ .where(pipeline_deployment_table.c.stack_id == stack_table.c.id)
110
+ .where(stack_composition_table.c.stack_id == stack_table.c.id)
111
+ .where(
112
+ stack_composition_table.c.component_id
113
+ == stack_component_table.c.id
114
+ )
115
+ .where(stack_component_table.c.flavor != "local")
116
+ .where(stack_component_table.c.type == "artifact_store")
117
+ ).scalar()
118
+ if (
119
+ pipeline_run_with_remote_artifact_store_count
120
+ and pipeline_run_with_remote_artifact_store_count > 0
121
+ ):
122
+ new_state.append("production_setup_completed")
123
+
124
+ pipeline_run_with_remote_orchestrator_count = connection.execute(
125
+ sa.select(sa.func.count(pipeline_run_table.c.id))
126
+ .where(
127
+ pipeline_run_table.c.deployment_id
128
+ == pipeline_deployment_table.c.id
129
+ )
130
+ .where(pipeline_deployment_table.c.stack_id == stack_table.c.id)
131
+ .where(stack_composition_table.c.stack_id == stack_table.c.id)
132
+ .where(
133
+ stack_composition_table.c.component_id
134
+ == stack_component_table.c.id
135
+ )
136
+ .where(
137
+ stack_component_table.c.flavor.not_in(["local", "local_docker"])
138
+ )
139
+ .where(stack_component_table.c.type == "orchestrator")
140
+ ).scalar()
141
+ if (
142
+ pipeline_run_with_remote_orchestrator_count
143
+ and pipeline_run_with_remote_orchestrator_count > 0
144
+ ):
145
+ new_state.append("pipeline_run_with_remote_orchestrator")
146
+ new_state.append("production_setup_completed")
147
+
148
+ if new_state:
149
+ # If any of the items are finished, we also complete the initial
150
+ # onboarding step which is not explicitly tracked in the database
151
+ new_state.append("device_verified")
152
+
153
+ # Remove duplicate keys
154
+ new_state = list(set(new_state))
155
+
156
+ connection.execute(
157
+ sa.update(server_settings_table).values(
158
+ onboarding_state=json.dumps(new_state)
159
+ )
160
+ )
161
+
162
+
163
+ def downgrade() -> None:
164
+ """Downgrade database schema and/or data back to the previous revision."""
165
+ # ### commands auto generated by Alembic - please adjust! ###
166
+ pass
167
+ # ### end Alembic commands ###
@@ -15,6 +15,7 @@
15
15
 
16
16
  import os
17
17
  import re
18
+ from datetime import datetime
18
19
  from pathlib import Path
19
20
  from typing import (
20
21
  Any,
@@ -57,6 +58,7 @@ from zenml.constants import (
57
58
  ARTIFACTS,
58
59
  CODE_REFERENCES,
59
60
  CODE_REPOSITORIES,
61
+ CONFIG,
60
62
  CURRENT_USER,
61
63
  DEACTIVATE,
62
64
  DEFAULT_HTTP_TIMEOUT,
@@ -65,6 +67,7 @@ from zenml.constants import (
65
67
  ENV_ZENML_DISABLE_CLIENT_SERVER_MISMATCH_WARNING,
66
68
  EVENT_SOURCES,
67
69
  FLAVORS,
70
+ FULL_STACK,
68
71
  GET_OR_CREATE,
69
72
  INFO,
70
73
  LOGIN,
@@ -89,9 +92,12 @@ from zenml.constants import (
89
92
  SERVICE_CONNECTOR_RESOURCES,
90
93
  SERVICE_CONNECTOR_TYPES,
91
94
  SERVICE_CONNECTOR_VERIFY,
95
+ SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT,
92
96
  SERVICE_CONNECTORS,
93
97
  SERVICES,
98
+ STACK,
94
99
  STACK_COMPONENTS,
100
+ STACK_DEPLOYMENT,
95
101
  STACKS,
96
102
  STEPS,
97
103
  TAGS,
@@ -103,6 +109,7 @@ from zenml.constants import (
103
109
  )
104
110
  from zenml.enums import (
105
111
  OAuthGrantTypes,
112
+ StackDeploymentProvider,
106
113
  StoreType,
107
114
  )
108
115
  from zenml.exceptions import AuthorizationException, MethodNotAllowedError
@@ -139,6 +146,7 @@ from zenml.models import (
139
146
  ComponentRequest,
140
147
  ComponentResponse,
141
148
  ComponentUpdate,
149
+ DeployedStack,
142
150
  EventSourceFilter,
143
151
  EventSourceRequest,
144
152
  EventSourceResponse,
@@ -147,6 +155,7 @@ from zenml.models import (
147
155
  FlavorRequest,
148
156
  FlavorResponse,
149
157
  FlavorUpdate,
158
+ FullStackRequest,
150
159
  LogsResponse,
151
160
  ModelFilter,
152
161
  ModelRequest,
@@ -208,6 +217,8 @@ from zenml.models import (
208
217
  ServiceRequest,
209
218
  ServiceResponse,
210
219
  ServiceUpdate,
220
+ StackDeploymentConfig,
221
+ StackDeploymentInfo,
211
222
  StackFilter,
212
223
  StackRequest,
213
224
  StackResponse,
@@ -2467,6 +2478,10 @@ class RestZenStore(BaseZenStore):
2467
2478
  f"{SERVICE_CONNECTORS}{SERVICE_CONNECTOR_VERIFY}",
2468
2479
  body=service_connector,
2469
2480
  params={"list_resources": list_resources},
2481
+ timeout=max(
2482
+ self.config.http_timeout,
2483
+ SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT,
2484
+ ),
2470
2485
  )
2471
2486
 
2472
2487
  resources = ServiceConnectorResourcesModel.model_validate(
@@ -2504,6 +2519,10 @@ class RestZenStore(BaseZenStore):
2504
2519
  response_body = self.put(
2505
2520
  f"{SERVICE_CONNECTORS}/{str(service_connector_id)}{SERVICE_CONNECTOR_VERIFY}",
2506
2521
  params=params,
2522
+ timeout=max(
2523
+ self.config.http_timeout,
2524
+ SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT,
2525
+ ),
2507
2526
  )
2508
2527
 
2509
2528
  resources = ServiceConnectorResourcesModel.model_validate(
@@ -2572,6 +2591,10 @@ class RestZenStore(BaseZenStore):
2572
2591
  response_body = self.get(
2573
2592
  f"{WORKSPACES}/{workspace_name_or_id}{SERVICE_CONNECTORS}{SERVICE_CONNECTOR_RESOURCES}",
2574
2593
  params=params,
2594
+ timeout=max(
2595
+ self.config.http_timeout,
2596
+ SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT,
2597
+ ),
2575
2598
  )
2576
2599
 
2577
2600
  assert isinstance(response_body, list)
@@ -2744,6 +2767,23 @@ class RestZenStore(BaseZenStore):
2744
2767
  response_model=StackResponse,
2745
2768
  )
2746
2769
 
2770
+ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse:
2771
+ """Register a full-stack.
2772
+
2773
+ Args:
2774
+ full_stack: The full stack configuration.
2775
+
2776
+ Returns:
2777
+ The registered stack.
2778
+ """
2779
+ assert full_stack.workspace is not None
2780
+
2781
+ return self._create_resource(
2782
+ resource=full_stack,
2783
+ response_model=StackResponse,
2784
+ route=f"{WORKSPACES}/{str(full_stack.workspace)}{FULL_STACK}",
2785
+ )
2786
+
2747
2787
  def get_stack(self, stack_id: UUID, hydrate: bool = True) -> StackResponse:
2748
2788
  """Get a stack by its unique ID.
2749
2789
 
@@ -2813,6 +2853,88 @@ class RestZenStore(BaseZenStore):
2813
2853
  route=STACKS,
2814
2854
  )
2815
2855
 
2856
+ # ---------------- Stack deployments-----------------
2857
+
2858
+ def get_stack_deployment_info(
2859
+ self,
2860
+ provider: StackDeploymentProvider,
2861
+ ) -> StackDeploymentInfo:
2862
+ """Get information about a stack deployment provider.
2863
+
2864
+ Args:
2865
+ provider: The stack deployment provider.
2866
+
2867
+ Returns:
2868
+ Information about the stack deployment provider.
2869
+ """
2870
+ body = self.get(
2871
+ f"{STACK_DEPLOYMENT}{INFO}",
2872
+ params={"provider": provider.value},
2873
+ )
2874
+ return StackDeploymentInfo.model_validate(body)
2875
+
2876
+ def get_stack_deployment_config(
2877
+ self,
2878
+ provider: StackDeploymentProvider,
2879
+ stack_name: str,
2880
+ location: Optional[str] = None,
2881
+ ) -> StackDeploymentConfig:
2882
+ """Return the cloud provider console URL and configuration needed to deploy the ZenML stack.
2883
+
2884
+ Args:
2885
+ provider: The stack deployment provider.
2886
+ stack_name: The name of the stack.
2887
+ location: The location where the stack should be deployed.
2888
+
2889
+ Returns:
2890
+ The cloud provider console URL and configuration needed to deploy
2891
+ the ZenML stack to the specified cloud provider.
2892
+ """
2893
+ params = {
2894
+ "provider": provider.value,
2895
+ "stack_name": stack_name,
2896
+ }
2897
+ if location:
2898
+ params["location"] = location
2899
+ body = self.get(f"{STACK_DEPLOYMENT}{CONFIG}", params=params)
2900
+ return StackDeploymentConfig.model_validate(body)
2901
+
2902
+ def get_stack_deployment_stack(
2903
+ self,
2904
+ provider: StackDeploymentProvider,
2905
+ stack_name: str,
2906
+ location: Optional[str] = None,
2907
+ date_start: Optional[datetime] = None,
2908
+ ) -> Optional[DeployedStack]:
2909
+ """Return a matching ZenML stack that was deployed and registered.
2910
+
2911
+ Args:
2912
+ provider: The stack deployment provider.
2913
+ stack_name: The name of the stack.
2914
+ location: The location where the stack should be deployed.
2915
+ date_start: The date when the deployment started.
2916
+
2917
+ Returns:
2918
+ The ZenML stack that was deployed and registered or None if the
2919
+ stack was not found.
2920
+ """
2921
+ params = {
2922
+ "provider": provider.value,
2923
+ "stack_name": stack_name,
2924
+ }
2925
+ if location:
2926
+ params["location"] = location
2927
+ if date_start:
2928
+ params["date_start"] = str(date_start)
2929
+ body = self.get(
2930
+ f"{STACK_DEPLOYMENT}{STACK}",
2931
+ params=params,
2932
+ )
2933
+ if body:
2934
+ return DeployedStack.model_validate(body)
2935
+
2936
+ return None
2937
+
2816
2938
  # ----------------------------- Step runs -----------------------------
2817
2939
 
2818
2940
  def create_run_step(self, step_run: StepRunRequest) -> StepRunResponse:
@@ -3964,6 +4086,7 @@ class RestZenStore(BaseZenStore):
3964
4086
  method: str,
3965
4087
  url: str,
3966
4088
  params: Optional[Dict[str, Any]] = None,
4089
+ timeout: Optional[int] = None,
3967
4090
  **kwargs: Any,
3968
4091
  ) -> Json:
3969
4092
  """Make a request to the REST API.
@@ -3972,6 +4095,7 @@ class RestZenStore(BaseZenStore):
3972
4095
  method: The HTTP method to use.
3973
4096
  url: The URL to request.
3974
4097
  params: The query parameters to pass to the endpoint.
4098
+ timeout: The request timeout in seconds.
3975
4099
  kwargs: Additional keyword arguments to pass to the request.
3976
4100
 
3977
4101
  Returns:
@@ -3994,7 +4118,7 @@ class RestZenStore(BaseZenStore):
3994
4118
  url,
3995
4119
  params=params,
3996
4120
  verify=self.config.verify_ssl,
3997
- timeout=self.config.http_timeout,
4121
+ timeout=timeout or self.config.http_timeout,
3998
4122
  **kwargs,
3999
4123
  )
4000
4124
  )
@@ -4023,13 +4147,18 @@ class RestZenStore(BaseZenStore):
4023
4147
  raise
4024
4148
 
4025
4149
  def get(
4026
- self, path: str, params: Optional[Dict[str, Any]] = None, **kwargs: Any
4150
+ self,
4151
+ path: str,
4152
+ params: Optional[Dict[str, Any]] = None,
4153
+ timeout: Optional[int] = None,
4154
+ **kwargs: Any,
4027
4155
  ) -> Json:
4028
4156
  """Make a GET request to the given endpoint path.
4029
4157
 
4030
4158
  Args:
4031
4159
  path: The path to the endpoint.
4032
4160
  params: The query parameters to pass to the endpoint.
4161
+ timeout: The request timeout in seconds.
4033
4162
  kwargs: Additional keyword arguments to pass to the request.
4034
4163
 
4035
4164
  Returns:
@@ -4037,17 +4166,26 @@ class RestZenStore(BaseZenStore):
4037
4166
  """
4038
4167
  logger.debug(f"Sending GET request to {path}...")
4039
4168
  return self._request(
4040
- "GET", self.url + API + VERSION_1 + path, params=params, **kwargs
4169
+ "GET",
4170
+ self.url + API + VERSION_1 + path,
4171
+ params=params,
4172
+ timeout=timeout,
4173
+ **kwargs,
4041
4174
  )
4042
4175
 
4043
4176
  def delete(
4044
- self, path: str, params: Optional[Dict[str, Any]] = None, **kwargs: Any
4177
+ self,
4178
+ path: str,
4179
+ params: Optional[Dict[str, Any]] = None,
4180
+ timeout: Optional[int] = None,
4181
+ **kwargs: Any,
4045
4182
  ) -> Json:
4046
4183
  """Make a DELETE request to the given endpoint path.
4047
4184
 
4048
4185
  Args:
4049
4186
  path: The path to the endpoint.
4050
4187
  params: The query parameters to pass to the endpoint.
4188
+ timeout: The request timeout in seconds.
4051
4189
  kwargs: Additional keyword arguments to pass to the request.
4052
4190
 
4053
4191
  Returns:
@@ -4058,6 +4196,7 @@ class RestZenStore(BaseZenStore):
4058
4196
  "DELETE",
4059
4197
  self.url + API + VERSION_1 + path,
4060
4198
  params=params,
4199
+ timeout=timeout,
4061
4200
  **kwargs,
4062
4201
  )
4063
4202
 
@@ -4066,6 +4205,7 @@ class RestZenStore(BaseZenStore):
4066
4205
  path: str,
4067
4206
  body: BaseModel,
4068
4207
  params: Optional[Dict[str, Any]] = None,
4208
+ timeout: Optional[int] = None,
4069
4209
  **kwargs: Any,
4070
4210
  ) -> Json:
4071
4211
  """Make a POST request to the given endpoint path.
@@ -4074,6 +4214,7 @@ class RestZenStore(BaseZenStore):
4074
4214
  path: The path to the endpoint.
4075
4215
  body: The body to send.
4076
4216
  params: The query parameters to pass to the endpoint.
4217
+ timeout: The request timeout in seconds.
4077
4218
  kwargs: Additional keyword arguments to pass to the request.
4078
4219
 
4079
4220
  Returns:
@@ -4085,6 +4226,7 @@ class RestZenStore(BaseZenStore):
4085
4226
  self.url + API + VERSION_1 + path,
4086
4227
  data=body.model_dump_json(),
4087
4228
  params=params,
4229
+ timeout=timeout,
4088
4230
  **kwargs,
4089
4231
  )
4090
4232
 
@@ -4093,6 +4235,7 @@ class RestZenStore(BaseZenStore):
4093
4235
  path: str,
4094
4236
  body: Optional[BaseModel] = None,
4095
4237
  params: Optional[Dict[str, Any]] = None,
4238
+ timeout: Optional[int] = None,
4096
4239
  **kwargs: Any,
4097
4240
  ) -> Json:
4098
4241
  """Make a PUT request to the given endpoint path.
@@ -4101,6 +4244,7 @@ class RestZenStore(BaseZenStore):
4101
4244
  path: The path to the endpoint.
4102
4245
  body: The body to send.
4103
4246
  params: The query parameters to pass to the endpoint.
4247
+ timeout: The request timeout in seconds.
4104
4248
  kwargs: Additional keyword arguments to pass to the request.
4105
4249
 
4106
4250
  Returns:
@@ -4113,6 +4257,7 @@ class RestZenStore(BaseZenStore):
4113
4257
  self.url + API + VERSION_1 + path,
4114
4258
  data=data,
4115
4259
  params=params,
4260
+ timeout=timeout,
4116
4261
  **kwargs,
4117
4262
  )
4118
4263
 
@@ -38,6 +38,7 @@ from zenml.zen_stores.schemas.user_schemas import UserSchema
38
38
  from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
39
39
 
40
40
  if TYPE_CHECKING:
41
+ from zenml.zen_stores.schemas.flavor_schemas import FlavorSchema
41
42
  from zenml.zen_stores.schemas.logs_schemas import LogsSchema
42
43
  from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
43
44
  from zenml.zen_stores.schemas.schedule_schema import ScheduleSchema
@@ -85,6 +86,13 @@ class StackComponentSchema(NamedSchema, table=True):
85
86
  run_metadata: List["RunMetadataSchema"] = Relationship(
86
87
  back_populates="stack_component",
87
88
  )
89
+ flavor_schema: Optional["FlavorSchema"] = Relationship(
90
+ sa_relationship_kwargs={
91
+ "primaryjoin": "and_(foreign(StackComponentSchema.flavor) == FlavorSchema.name, foreign(StackComponentSchema.type) == FlavorSchema.type)",
92
+ "lazy": "joined",
93
+ "uselist": False,
94
+ },
95
+ )
88
96
 
89
97
  run_or_step_logs: List["LogsSchema"] = Relationship(
90
98
  back_populates="artifact_store",
@@ -161,6 +169,12 @@ class StackComponentSchema(NamedSchema, table=True):
161
169
  user=self.user.to_model() if self.user else None,
162
170
  created=self.created,
163
171
  updated=self.updated,
172
+ logo_url=self.flavor_schema.logo_url
173
+ if self.flavor_schema
174
+ else None,
175
+ integration=self.flavor_schema.integration
176
+ if self.flavor_schema
177
+ else None,
164
178
  )
165
179
  metadata = None
166
180
  if include_metadata: