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
@@ -0,0 +1,596 @@
1
+ # Copyright (c) ZenML GmbH 2023. 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
+ """Implementation of the Lightning orchestrator."""
15
+
16
+ import os
17
+ import tempfile
18
+ import time
19
+ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, cast
20
+ from uuid import uuid4
21
+
22
+ from lightning_sdk import Machine, Studio
23
+
24
+ from zenml.constants import (
25
+ ENV_ZENML_CUSTOM_SOURCE_ROOT,
26
+ ENV_ZENML_WHEEL_PACKAGE_NAME,
27
+ )
28
+ from zenml.entrypoints.step_entrypoint_configuration import (
29
+ StepEntrypointConfiguration,
30
+ )
31
+ from zenml.integrations.lightning.flavors.lightning_orchestrator_flavor import (
32
+ LightningOrchestratorConfig,
33
+ LightningOrchestratorSettings,
34
+ )
35
+ from zenml.integrations.lightning.orchestrators.lightning_orchestrator_entrypoint_configuration import (
36
+ LightningOrchestratorEntrypointConfiguration,
37
+ )
38
+ from zenml.integrations.lightning.orchestrators.utils import (
39
+ gather_requirements,
40
+ sanitize_studio_name,
41
+ )
42
+ from zenml.logger import get_logger
43
+ from zenml.orchestrators.utils import get_orchestrator_run_name
44
+ from zenml.orchestrators.wheeled_orchestrator import WheeledOrchestrator
45
+ from zenml.stack import StackValidator
46
+ from zenml.utils import code_utils, io_utils, source_utils
47
+
48
+ if TYPE_CHECKING:
49
+ from zenml.models import PipelineDeploymentResponse
50
+ from zenml.stack import Stack
51
+
52
+
53
+ logger = get_logger(__name__)
54
+ ENV_ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID = "ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID"
55
+ ZENML_STEP_DEFAULT_ENTRYPOINT_COMMAND = "zenml.entrypoints.entrypoint"
56
+ LIGHTNING_ZENML_DEFAULT_CUSTOM_REPOSITORY_PATH = "."
57
+
58
+
59
+ class LightningOrchestrator(WheeledOrchestrator):
60
+ """Base class for Orchestrator responsible for running pipelines remotely in a VM.
61
+
62
+ This orchestrator does not support running on a schedule.
63
+ """
64
+
65
+ @property
66
+ def validator(self) -> Optional[StackValidator]:
67
+ """Validates the stack.
68
+
69
+ In the remote case, checks that the stack contains a container registry,
70
+ image builder and only remote components.
71
+
72
+ Returns:
73
+ A `StackValidator` instance.
74
+ """
75
+
76
+ def _validate_remote_components(
77
+ stack: "Stack",
78
+ ) -> Tuple[bool, str]:
79
+ for component in stack.components.values():
80
+ if not component.config.is_local:
81
+ continue
82
+
83
+ # return False, (
84
+ # f"The Lightning orchestrator runs pipelines remotely, "
85
+ # f"but the '{component.name}' {component.type.value} is "
86
+ # "a local stack component and will not be available in "
87
+ # "the Lightning step.\nPlease ensure that you always "
88
+ # "use non-local stack components with the Lightning "
89
+ # "orchestrator."
90
+ # )
91
+
92
+ return True, ""
93
+
94
+ return StackValidator(
95
+ custom_validation_function=_validate_remote_components,
96
+ )
97
+
98
+ def _set_lightning_env_vars(
99
+ self,
100
+ deployment: "PipelineDeploymentResponse",
101
+ ) -> None:
102
+ """Set up the Lightning client using environment variables.
103
+
104
+ Args:
105
+ deployment: The pipeline deployment to prepare or run.
106
+ """
107
+ settings = cast(
108
+ LightningOrchestratorSettings, self.get_settings(deployment)
109
+ )
110
+ if settings.user_id:
111
+ os.environ["LIGHTNING_USER_ID"] = settings.user_id
112
+ if settings.api_key:
113
+ os.environ["LIGHTNING_API_KEY"] = settings.api_key
114
+ if settings.username:
115
+ os.environ["LIGHTNING_USERNAME"] = settings.username
116
+ if settings.teamspace:
117
+ os.environ["LIGHTNING_TEAMSPACE"] = settings.teamspace
118
+ if settings.organization:
119
+ os.environ["LIGHTNING_ORG"] = settings.organization
120
+
121
+ @property
122
+ def config(self) -> LightningOrchestratorConfig:
123
+ """Returns the `LightningOrchestratorConfig` config.
124
+
125
+ Returns:
126
+ The configuration.
127
+ """
128
+ return cast(LightningOrchestratorConfig, self._config)
129
+
130
+ @property
131
+ def settings_class(self) -> Type[LightningOrchestratorSettings]:
132
+ """Settings class for the Lightning orchestrator.
133
+
134
+ Returns:
135
+ The settings class.
136
+ """
137
+ return LightningOrchestratorSettings
138
+
139
+ def get_orchestrator_run_id(self) -> str:
140
+ """Returns the active orchestrator run id.
141
+
142
+ Raises:
143
+ RuntimeError: If no run id exists. This happens when this method
144
+ gets called while the orchestrator is not running a pipeline.
145
+
146
+ Returns:
147
+ The orchestrator run id.
148
+
149
+ Raises:
150
+ RuntimeError: If the run id cannot be read from the environment.
151
+ """
152
+ try:
153
+ return os.environ[ENV_ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID]
154
+ except KeyError:
155
+ raise RuntimeError(
156
+ "Unable to read run id from environment variable "
157
+ f"{ENV_ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID}."
158
+ )
159
+
160
+ @property
161
+ def root_directory(self) -> str:
162
+ """Path to the root directory for all files concerning this orchestrator.
163
+
164
+ Returns:
165
+ Path to the root directory.
166
+ """
167
+ return os.path.join(
168
+ io_utils.get_global_config_directory(),
169
+ "lightning",
170
+ str(self.id),
171
+ )
172
+
173
+ @property
174
+ def pipeline_directory(self) -> str:
175
+ """Returns path to a directory in which the kubeflow pipeline files are stored.
176
+
177
+ Returns:
178
+ Path to the pipeline directory.
179
+ """
180
+ return os.path.join(self.root_directory, "pipelines")
181
+
182
+ def setup_credentials(self) -> None:
183
+ """Set up credentials for the orchestrator."""
184
+ connector = self.get_connector()
185
+ assert connector is not None
186
+ connector.configure_local_client()
187
+
188
+ def prepare_or_run_pipeline(
189
+ self,
190
+ deployment: "PipelineDeploymentResponse",
191
+ stack: "Stack",
192
+ environment: Dict[str, str],
193
+ ) -> Any:
194
+ """Creates a wheel and uploads the pipeline to Lightning.
195
+
196
+ This functions as an intermediary representation of the pipeline which
197
+ is then deployed to the kubeflow pipelines instance.
198
+
199
+ How it works:
200
+ -------------
201
+ Before this method is called the `prepare_pipeline_deployment()`
202
+ method builds a docker image that contains the code for the
203
+ pipeline, all steps the context around these files.
204
+
205
+ Based on this docker image a callable is created which builds
206
+ task for each step (`_construct_lightning_pipeline`).
207
+ To do this the entrypoint of the docker image is configured to
208
+ run the correct step within the docker image. The dependencies
209
+ between these task are then also configured onto each
210
+ task by pointing at the downstream steps.
211
+
212
+ Args:
213
+ deployment: The pipeline deployment to prepare or run.
214
+ stack: The stack the pipeline will run on.
215
+ environment: Environment variables to set in the orchestration
216
+ environment.
217
+
218
+ Raises:
219
+ ValueError: If the schedule is not set or if the cron expression
220
+ is not set.
221
+ """
222
+ settings = cast(
223
+ LightningOrchestratorSettings, self.get_settings(deployment)
224
+ )
225
+ if deployment.schedule:
226
+ if (
227
+ deployment.schedule.catchup
228
+ or deployment.schedule.interval_second
229
+ ):
230
+ logger.warning(
231
+ "Lightning orchestrator only uses schedules with the "
232
+ "`cron_expression` property, with optional `start_time` and/or `end_time`. "
233
+ "All other properties are ignored."
234
+ )
235
+ if deployment.schedule.cron_expression is None:
236
+ raise ValueError(
237
+ "Property `cron_expression` must be set when passing "
238
+ "schedule to a Lightning orchestrator."
239
+ )
240
+ if deployment.schedule.cron_expression:
241
+ raise ValueError(
242
+ "Property `schedule_timezone` must be set when passing "
243
+ "`cron_expression` to a Lightning orchestrator."
244
+ "Lightning orchestrator requires a Java Timezone ID to run the pipeline on schedule."
245
+ "Please refer to https://docs.oracle.com/middleware/1221/wcs/tag-ref/MISC/TimeZones.html for more information."
246
+ )
247
+
248
+ # Get deployment id
249
+ deployment_id = deployment.id
250
+
251
+ pipeline_name = deployment.pipeline_configuration.name
252
+ orchestrator_run_name = get_orchestrator_run_name(pipeline_name)
253
+
254
+ # Copy the repository to a temporary directory and add a setup.py file
255
+ # repository_temp_dir = (
256
+ # self.copy_repository_to_temp_dir_and_add_setup_py()
257
+ # )
258
+
259
+ # Create a wheel for the package in the temporary directory
260
+ # wheel_path = self.create_wheel(temp_dir=repository_temp_dir)
261
+ code_archive = code_utils.CodeArchive(
262
+ root=source_utils.get_source_root()
263
+ )
264
+ logger.info("Archiving pipeline code...")
265
+ with tempfile.NamedTemporaryFile(
266
+ mode="w+b", delete=False, suffix=".tar.gz"
267
+ ) as code_file:
268
+ code_archive.write_archive(code_file)
269
+ code_path = code_file.name
270
+
271
+ filename = f"{orchestrator_run_name}.tar.gz"
272
+
273
+ # Construct the env variables for the pipeline
274
+ env_vars = environment.copy()
275
+ orchestrator_run_id = str(uuid4())
276
+ env_vars[ENV_ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID] = orchestrator_run_id
277
+ # Set up some variables for configuration
278
+ env_vars[ENV_ZENML_CUSTOM_SOURCE_ROOT] = (
279
+ LIGHTNING_ZENML_DEFAULT_CUSTOM_REPOSITORY_PATH
280
+ )
281
+ env_vars[ENV_ZENML_WHEEL_PACKAGE_NAME] = self.package_name
282
+
283
+ # Create a line-by-line export of environment variables
284
+ env_exports = "\n".join(
285
+ [f"export {key}='{value}'" for key, value in env_vars.items()]
286
+ )
287
+
288
+ # Write the environment variables to a temporary file
289
+ with tempfile.NamedTemporaryFile(
290
+ mode="w", delete=False, suffix=".studiorc"
291
+ ) as temp_file:
292
+ temp_file.write(env_exports)
293
+ env_file_path = temp_file.name
294
+
295
+ # Gather the requirements
296
+ pipeline_docker_settings = (
297
+ deployment.pipeline_configuration.docker_settings
298
+ )
299
+ pipeline_requirements = gather_requirements(pipeline_docker_settings)
300
+ pipeline_requirements_to_string = " ".join(
301
+ f'"{req}"' for req in pipeline_requirements
302
+ )
303
+
304
+ def _construct_lightning_steps(
305
+ deployment: "PipelineDeploymentResponse",
306
+ ) -> Dict[str, Dict[str, Any]]:
307
+ """Construct the steps for the pipeline.
308
+
309
+ Args:
310
+ deployment: The pipeline deployment to prepare or run.
311
+
312
+ Returns:
313
+ The steps for the pipeline.
314
+ """
315
+ steps = {}
316
+ for step_name, step in deployment.step_configurations.items():
317
+ # The arguments are passed to configure the entrypoint of the
318
+ # docker container when the step is called.
319
+ entrypoint_command = (
320
+ StepEntrypointConfiguration.get_entrypoint_command()
321
+ )
322
+ entrypoint_arguments = (
323
+ StepEntrypointConfiguration.get_entrypoint_arguments(
324
+ step_name=step_name,
325
+ deployment_id=deployment_id,
326
+ )
327
+ )
328
+ entrypoint = entrypoint_command + entrypoint_arguments
329
+ entrypoint_string = " ".join(entrypoint)
330
+
331
+ step_settings = cast(
332
+ LightningOrchestratorSettings, self.get_settings(step)
333
+ )
334
+
335
+ # Gather the requirements
336
+ step_docker_settings = step.config.docker_settings
337
+ step_requirements = gather_requirements(step_docker_settings)
338
+ step_requirements_to_string = " ".join(
339
+ f'"{req}"' for req in step_requirements
340
+ )
341
+
342
+ # Construct the command to run the step
343
+ run_command = f"{entrypoint_string}"
344
+ commands = [run_command]
345
+ steps[step_name] = {
346
+ "commands": commands,
347
+ "requirements": step_requirements_to_string,
348
+ "machine": step_settings.machine_type
349
+ if step_settings != settings
350
+ else None,
351
+ }
352
+ return steps
353
+
354
+ if not settings.synchronous:
355
+ entrypoint_command = LightningOrchestratorEntrypointConfiguration.get_entrypoint_command()
356
+ entrypoint_arguments = LightningOrchestratorEntrypointConfiguration.get_entrypoint_arguments(
357
+ run_name=orchestrator_run_name,
358
+ deployment_id=deployment.id,
359
+ )
360
+ entrypoint = entrypoint_command + entrypoint_arguments
361
+ entrypoint_string = " ".join(entrypoint)
362
+ logger.info("Setting up Lightning AI client")
363
+ self._set_lightning_env_vars(deployment)
364
+
365
+ studio_name = sanitize_studio_name(
366
+ "zenml_async_orchestrator_studio"
367
+ )
368
+ logger.info(f"Creating main studio: {studio_name}")
369
+ studio = Studio(name=studio_name)
370
+ studio.start()
371
+
372
+ logger.info(
373
+ "Uploading wheel package and installing dependencies on main studio"
374
+ )
375
+ studio.run(
376
+ f"mkdir -p /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
377
+ )
378
+ studio.upload_file(
379
+ code_path,
380
+ remote_path=f"/teamspace/studios/this_studio/zenml_codes/{filename}",
381
+ )
382
+ time.sleep(10)
383
+ studio.run(
384
+ f"tar -xvzf /teamspace/studios/this_studio/zenml_codes/{filename} -C /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
385
+ )
386
+ studio.upload_file(
387
+ env_file_path,
388
+ remote_path="/teamspace/studios/this_studio/.lightning_studio/.studiorc",
389
+ )
390
+ studio.run("pip install uv")
391
+ logger.info(
392
+ f"Installing requirements: {pipeline_requirements_to_string}"
393
+ )
394
+ studio.run(f"uv pip install {pipeline_requirements_to_string}")
395
+ studio.run(
396
+ "pip uninstall zenml -y && pip install git+https://github.com/zenml-io/zenml.git@feature/lightening-studio-orchestrator"
397
+ )
398
+
399
+ for custom_command in settings.custom_commands or []:
400
+ studio.run(
401
+ f"cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {custom_command}"
402
+ )
403
+ # studio.run(f"pip install {wheel_path.rsplit('/', 1)[-1]}")
404
+ logger.info("Running pipeline in async mode")
405
+ studio.run(
406
+ f"nohup bash -c 'cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {entrypoint_string}' > log_{filename.rsplit('.', 2)[0]}.txt 2>&1 &"
407
+ )
408
+ logger.info(
409
+ f"The pipeline is running in async mode, you can keep checking the logs by running the following command: `lightning download -s vision-model/zenml-async-orchestrator-studio -p /teamspace/studios/this_studio/log_{filename.rsplit('.', 2)[0]}.txt && cat log_{filename.rsplit('.', 2)[0]}.txt`"
410
+ )
411
+ else:
412
+ self._upload_and_run_pipeline(
413
+ deployment,
414
+ orchestrator_run_id,
415
+ pipeline_requirements_to_string,
416
+ settings,
417
+ _construct_lightning_steps(deployment),
418
+ code_path,
419
+ filename,
420
+ env_file_path,
421
+ )
422
+ os.unlink(env_file_path)
423
+
424
+ def _upload_and_run_pipeline(
425
+ self,
426
+ deployment: "PipelineDeploymentResponse",
427
+ orchestrator_run_id: str,
428
+ requirements: str,
429
+ settings: LightningOrchestratorSettings,
430
+ steps_commands: Dict[str, Dict[str, Any]],
431
+ code_path: str,
432
+ filename: str,
433
+ env_file_path: str,
434
+ ) -> None:
435
+ """Upload and run the pipeline on Lightning Studio.
436
+
437
+ Args:
438
+ deployment: The pipeline deployment to prepare or run.
439
+ orchestrator_run_id: The orchestrator run id.
440
+ requirements: The requirements for the pipeline.
441
+ settings: The orchestrator settings.
442
+ steps_commands: The commands for the steps.
443
+ code_path: The path to the wheel package.
444
+ filename: The name of the code archive.
445
+ env_file_path: The path to the environment file.
446
+
447
+ Raises:
448
+ Exception: If an error occurs while running the pipeline.
449
+ """
450
+ logger.info("Setting up Lightning AI client")
451
+ self._set_lightning_env_vars(deployment)
452
+
453
+ if settings.main_studio_name:
454
+ studio_name = settings.main_studio_name
455
+ studio = Studio(name=studio_name)
456
+ if (
457
+ studio.machine != settings.machine_type
458
+ and settings.machine_type
459
+ ):
460
+ studio.switch_machine(Machine(settings.machine_type))
461
+ else:
462
+ studio_name = sanitize_studio_name(
463
+ f"zenml_{orchestrator_run_id}_pipeline"
464
+ )
465
+ logger.info(f"Creating main studio: {studio_name}")
466
+ studio = Studio(name=studio_name)
467
+ if settings.machine_type:
468
+ studio.start(Machine(settings.machine_type))
469
+ else:
470
+ studio.start()
471
+ try:
472
+ studio.run(
473
+ f"mkdir -p /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
474
+ )
475
+ studio.upload_file(
476
+ code_path,
477
+ remote_path=f"/teamspace/studios/this_studio/zenml_codes/{filename}",
478
+ )
479
+ studio.run(
480
+ f"tar -xvzf /teamspace/studios/this_studio/zenml_codes/{filename} -C /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
481
+ )
482
+ logger.info(
483
+ "Uploading wheel package and installing dependencies on main studio"
484
+ )
485
+ studio.upload_file(
486
+ env_file_path,
487
+ remote_path="/teamspace/studios/this_studio/.lightning_studio/.studiorc",
488
+ )
489
+ studio.run("pip install uv")
490
+ studio.run(f"uv pip install {requirements}")
491
+ studio.run(
492
+ "pip uninstall zenml -y && pip install git+https://github.com/zenml-io/zenml.git@feature/lightening-studio-orchestrator"
493
+ )
494
+ # studio.run(f"pip install {wheel_path.rsplit('/', 1)[-1]}")
495
+ for command in settings.custom_commands or []:
496
+ output = studio.run(
497
+ f"cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {command}"
498
+ )
499
+ logger.info(f"Custom command output: {output}")
500
+
501
+ for step_name, details in steps_commands.items():
502
+ if details["machine"]:
503
+ logger.info(f"Executing step: {step_name} in new studio")
504
+ self._run_step_in_new_studio(
505
+ orchestrator_run_id,
506
+ step_name,
507
+ details,
508
+ code_path,
509
+ filename,
510
+ env_file_path,
511
+ settings.custom_commands,
512
+ )
513
+ else:
514
+ logger.info(f"Executing step: {step_name} in main studio")
515
+ self._run_step_in_main_studio(studio, details, filename)
516
+ except Exception as e:
517
+ logger.error(f"Error running pipeline: {e}")
518
+ raise e
519
+ finally:
520
+ if (
521
+ studio.status != studio.status.NotCreated
522
+ and settings.main_studio_name is None
523
+ ):
524
+ logger.info("Deleting main studio")
525
+ studio.delete()
526
+
527
+ def _run_step_in_new_studio(
528
+ self,
529
+ orchestrator_run_id: str,
530
+ step_name: str,
531
+ details: Dict[str, Any],
532
+ code_path: str,
533
+ filename: str,
534
+ env_file_path: str,
535
+ custom_commands: Optional[List[str]] = None,
536
+ ) -> None:
537
+ """Run a step in a new studio.
538
+
539
+ Args:
540
+ orchestrator_run_id: The orchestrator run id.
541
+ step_name: The name of the step.
542
+ details: The details of the step.
543
+ code_path: The path to the wheel package.
544
+ filename: The name of the code archive.
545
+ env_file_path: The path to the environment file.
546
+ custom_commands: Custom commands to run.
547
+ """
548
+ studio_name = sanitize_studio_name(
549
+ f"zenml_{orchestrator_run_id}_{step_name}"
550
+ )
551
+ logger.info(f"Creating new studio for step {step_name}: {studio_name}")
552
+ studio = Studio(name=studio_name)
553
+ studio.start(Machine(details["machine"]))
554
+ studio.run(
555
+ f"mkdir -p /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
556
+ )
557
+ studio.upload_file(code_path, remote_path=f"/zenml_codes/{filename}")
558
+ studio.run(
559
+ f"tar -xvzf /teamspace/studios/this_studio/zenml_codes/{filename} -C /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
560
+ )
561
+ studio.upload_file(
562
+ env_file_path, remote_path=".lightning_studio/.studiorc"
563
+ )
564
+ studio.run("pip install uv")
565
+ studio.run(f"uv pip install {details['requirements']}")
566
+ studio.run(
567
+ "pip uninstall zenml -y && pip install git+https://github.com/zenml-io/zenml.git@feature/lightening-studio-orchestrator"
568
+ )
569
+ # studio.run(f"pip install {wheel_path.rsplit('/', 1)[-1]}")
570
+ for command in custom_commands or []:
571
+ output = studio.run(
572
+ f"cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {command}"
573
+ )
574
+ logger.info(f"Custom command output: {output}")
575
+ for command in details["commands"]:
576
+ output = studio.run(
577
+ f"cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {command}"
578
+ )
579
+ logger.info(f"Step {step_name} output: {output}")
580
+ studio.delete()
581
+
582
+ def _run_step_in_main_studio(
583
+ self, studio: Studio, details: Dict[str, Any], filename: str
584
+ ) -> None:
585
+ """Run a step in the main studio.
586
+
587
+ Args:
588
+ studio: The studio to run the step in.
589
+ details: The details of the step.
590
+ filename: The name of the code archive.
591
+ """
592
+ for command in details["commands"]:
593
+ output = studio.run(
594
+ f"cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {command}"
595
+ )
596
+ logger.info(f"Step output: {output}")