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.
- README.md +1 -1
- RELEASE_NOTES.md +126 -4
- zenml/VERSION +1 -1
- zenml/artifacts/utils.py +13 -6
- zenml/cli/__init__.py +1 -1
- zenml/cli/base.py +4 -4
- zenml/cli/integration.py +48 -9
- zenml/cli/pipeline.py +9 -2
- zenml/cli/stack.py +39 -27
- zenml/cli/utils.py +13 -0
- zenml/client.py +15 -17
- zenml/config/compiler.py +34 -0
- zenml/config/server_config.py +30 -0
- zenml/config/source.py +3 -7
- zenml/constants.py +5 -3
- zenml/entrypoints/base_entrypoint_configuration.py +41 -27
- zenml/entrypoints/step_entrypoint_configuration.py +5 -2
- zenml/enums.py +2 -0
- zenml/environment.py +31 -0
- zenml/feature_stores/base_feature_store.py +4 -6
- zenml/integrations/__init__.py +3 -0
- zenml/integrations/airflow/flavors/airflow_orchestrator_flavor.py +9 -0
- zenml/integrations/aws/__init__.py +2 -2
- zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py +2 -1
- zenml/integrations/azure/__init__.py +2 -2
- zenml/integrations/azure/azureml_utils.py +201 -0
- zenml/integrations/azure/flavors/azureml.py +139 -0
- zenml/integrations/azure/flavors/azureml_orchestrator_flavor.py +20 -118
- zenml/integrations/azure/flavors/azureml_step_operator_flavor.py +67 -14
- zenml/integrations/azure/orchestrators/azureml_orchestrator.py +58 -172
- zenml/integrations/azure/orchestrators/azureml_orchestrator_entrypoint_config.py +1 -0
- zenml/integrations/azure/service_connectors/azure_service_connector.py +4 -0
- zenml/integrations/azure/step_operators/azureml_step_operator.py +78 -177
- zenml/integrations/constants.py +3 -0
- zenml/integrations/databricks/__init__.py +22 -4
- zenml/integrations/databricks/flavors/databricks_orchestrator_flavor.py +9 -0
- zenml/integrations/deepchecks/__init__.py +29 -11
- zenml/integrations/deepchecks/materializers/deepchecks_dataset_materializer.py +3 -1
- zenml/integrations/deepchecks/validation_checks.py +0 -30
- zenml/integrations/evidently/__init__.py +17 -2
- zenml/integrations/facets/__init__.py +21 -5
- zenml/integrations/feast/__init__.py +19 -6
- zenml/integrations/gcp/__init__.py +2 -2
- zenml/integrations/gcp/flavors/vertex_orchestrator_flavor.py +9 -0
- zenml/integrations/gcp/orchestrators/vertex_orchestrator.py +10 -1
- zenml/integrations/great_expectations/__init__.py +21 -7
- zenml/integrations/huggingface/__init__.py +39 -15
- zenml/integrations/huggingface/materializers/__init__.py +3 -0
- zenml/integrations/huggingface/materializers/huggingface_datasets_materializer.py +3 -1
- zenml/integrations/huggingface/materializers/huggingface_pt_model_materializer.py +1 -1
- zenml/integrations/huggingface/materializers/huggingface_t5_materializer.py +107 -0
- zenml/integrations/huggingface/materializers/huggingface_tf_model_materializer.py +1 -1
- zenml/integrations/huggingface/materializers/huggingface_tokenizer_materializer.py +2 -2
- zenml/integrations/huggingface/steps/accelerate_runner.py +108 -85
- zenml/integrations/hyperai/flavors/hyperai_orchestrator_flavor.py +9 -0
- zenml/integrations/kubeflow/flavors/kubeflow_orchestrator_flavor.py +9 -0
- zenml/integrations/kubeflow/orchestrators/kubeflow_orchestrator.py +10 -1
- zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py +9 -0
- zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py +10 -1
- zenml/integrations/lightning/__init__.py +48 -0
- zenml/integrations/lightning/flavors/__init__.py +23 -0
- zenml/integrations/lightning/flavors/lightning_orchestrator_flavor.py +148 -0
- zenml/integrations/lightning/orchestrators/__init__.py +23 -0
- zenml/integrations/lightning/orchestrators/lightning_orchestrator.py +596 -0
- zenml/integrations/lightning/orchestrators/lightning_orchestrator_entrypoint.py +307 -0
- zenml/integrations/lightning/orchestrators/lightning_orchestrator_entrypoint_configuration.py +77 -0
- zenml/integrations/lightning/orchestrators/utils.py +67 -0
- zenml/integrations/mlflow/__init__.py +43 -5
- zenml/integrations/mlflow/services/mlflow_deployment.py +26 -0
- zenml/integrations/numpy/__init__.py +32 -0
- zenml/integrations/numpy/materializers/__init__.py +18 -0
- zenml/integrations/numpy/materializers/numpy_materializer.py +246 -0
- zenml/integrations/pandas/__init__.py +32 -0
- zenml/integrations/pandas/materializers/__init__.py +18 -0
- zenml/integrations/pandas/materializers/pandas_materializer.py +192 -0
- zenml/integrations/prodigy/annotators/prodigy_annotator.py +1 -1
- zenml/integrations/seldon/__init__.py +18 -3
- zenml/integrations/sklearn/__init__.py +1 -1
- zenml/integrations/skypilot_azure/__init__.py +1 -1
- zenml/integrations/tensorboard/__init__.py +1 -1
- zenml/integrations/tensorflow/__init__.py +2 -2
- zenml/integrations/tensorflow/materializers/tf_dataset_materializer.py +2 -2
- zenml/integrations/whylogs/__init__.py +18 -2
- zenml/logging/step_logging.py +9 -2
- zenml/materializers/__init__.py +0 -4
- zenml/materializers/base_materializer.py +4 -0
- zenml/materializers/numpy_materializer.py +23 -234
- zenml/materializers/pandas_materializer.py +22 -179
- zenml/model/model.py +91 -2
- zenml/model/utils.py +5 -5
- zenml/models/__init__.py +16 -3
- zenml/models/v2/core/model_version.py +1 -1
- zenml/models/v2/core/pipeline_run.py +31 -1
- zenml/models/v2/core/stack.py +51 -20
- zenml/models/v2/core/step_run.py +28 -0
- zenml/models/v2/misc/info_models.py +78 -0
- zenml/new/pipelines/pipeline.py +65 -25
- zenml/new/pipelines/run_utils.py +57 -136
- zenml/new/steps/step_context.py +17 -6
- zenml/orchestrators/base_orchestrator.py +9 -0
- zenml/orchestrators/step_launcher.py +37 -14
- zenml/orchestrators/step_runner.py +14 -13
- zenml/orchestrators/utils.py +107 -7
- zenml/service_connectors/service_connector_utils.py +2 -2
- zenml/stack/utils.py +11 -2
- zenml/stack_deployments/azure_stack_deployment.py +2 -1
- zenml/steps/base_step.py +62 -25
- zenml/steps/utils.py +115 -3
- zenml/utils/cloud_utils.py +8 -8
- zenml/utils/code_utils.py +130 -32
- zenml/utils/function_utils.py +7 -7
- zenml/utils/notebook_utils.py +14 -0
- zenml/utils/pipeline_docker_image_builder.py +1 -11
- zenml/utils/pydantic_utils.py +3 -3
- zenml/utils/secret_utils.py +2 -2
- zenml/utils/settings_utils.py +1 -1
- zenml/utils/source_utils.py +67 -21
- zenml/utils/string_utils.py +29 -0
- zenml/zen_server/dashboard/assets/{404-CRAA_Lew.js → 404-iO8vpun1.js} +1 -1
- zenml/zen_server/dashboard/assets/{@radix-BXWm7HOa.js → @radix-DnFH_oo1.js} +1 -1
- zenml/zen_server/dashboard/assets/{@react-router-l3lMcXA2.js → @react-router-APVeuk-U.js} +1 -1
- zenml/zen_server/dashboard/assets/{@reactflow-CeVxyqYT.js → @reactflow-B6kq9fJZ.js} +2 -2
- zenml/zen_server/dashboard/assets/{@tanstack-FmcYZMuX.js → @tanstack-QbMbTrh5.js} +1 -1
- zenml/zen_server/dashboard/assets/AlertDialogDropdownItem-BXeSvmMY.js +1 -0
- zenml/zen_server/dashboard/assets/{CodeSnippet-D0VLxT2A.js → CodeSnippet-DNWdQmbo.js} +2 -2
- zenml/zen_server/dashboard/assets/CollapsibleCard-B2OVjWYE.js +1 -0
- zenml/zen_server/dashboard/assets/Commands-DsoaVElZ.js +1 -0
- zenml/zen_server/dashboard/assets/CopyButton-BqE_-PHO.js +2 -0
- zenml/zen_server/dashboard/assets/{CsvVizualization-D3kAypDj.js → CsvVizualization-Dyasr2jU.js} +6 -6
- zenml/zen_server/dashboard/assets/{edit-C0MVvPD2.js → DialogItem-Cz1VLRwa.js} +1 -1
- zenml/zen_server/dashboard/assets/{DisplayDate-DizbSeT-.js → DisplayDate-DkCy54Bp.js} +1 -1
- zenml/zen_server/dashboard/assets/EditSecretDialog-Du423_3U.js +1 -0
- zenml/zen_server/dashboard/assets/{EmptyState-BHblM39I.js → EmptyState-Cs3DEmso.js} +1 -1
- zenml/zen_server/dashboard/assets/{Error-C6LeJSER.js → Error-DorJD_va.js} +1 -1
- zenml/zen_server/dashboard/assets/ExecutionStatus-CIfQTutR.js +1 -0
- zenml/zen_server/dashboard/assets/{Helpbox-aAB2XP-z.js → Helpbox-CmfvtNeq.js} +1 -1
- zenml/zen_server/dashboard/assets/Infobox-BL9NOS37.js +1 -0
- zenml/zen_server/dashboard/assets/{InlineAvatar-DpTLgM3Q.js → InlineAvatar-Ds2ZFHPc.js} +1 -1
- zenml/zen_server/dashboard/assets/{Lock-CNyJvf2r.js → Lock-CmIn0szs.js} +1 -1
- zenml/zen_server/dashboard/assets/{MarkdownVisualization-Bajxn0HY.js → MarkdownVisualization-DS05sfBm.js} +1 -1
- zenml/zen_server/dashboard/assets/{NumberBox-BmKE0qnO.js → NumberBox-CrN0_kqI.js} +1 -1
- zenml/zen_server/dashboard/assets/Partials-DX-8iEa1.js +1 -0
- zenml/zen_server/dashboard/assets/{PasswordChecker-yGGoJSB-.js → PasswordChecker-DE71J_3F.js} +1 -1
- zenml/zen_server/dashboard/assets/ProviderIcon-BOQJgapd.js +1 -0
- zenml/zen_server/dashboard/assets/ProviderRadio-BsYBw9YA.js +1 -0
- zenml/zen_server/dashboard/assets/SearchField-W3GXpLlI.js +1 -0
- zenml/zen_server/dashboard/assets/SetPassword-B-0a8UCj.js +1 -0
- zenml/zen_server/dashboard/assets/{Tick-uxv80Q6a.js → Tick-i1DYsVcX.js} +1 -1
- zenml/zen_server/dashboard/assets/{UpdatePasswordSchemas-oN4G3sKz.js → UpdatePasswordSchemas-C6Zb7ASL.js} +1 -1
- zenml/zen_server/dashboard/assets/UsageReason-CCnzmwS8.js +1 -0
- zenml/zen_server/dashboard/assets/WizardFooter-BHbO7zOa.js +1 -0
- zenml/zen_server/dashboard/assets/all-pipeline-runs-query-BBEe6I9-.js +1 -0
- zenml/zen_server/dashboard/assets/{check-circle-1_I207rW.js → check-circle-DOoS4yhF.js} +1 -1
- zenml/zen_server/dashboard/assets/{chevron-down-BpaF8JqM.js → chevron-down-Cwb-W_B_.js} +1 -1
- zenml/zen_server/dashboard/assets/{chevron-right-double-Dk8e2L99.js → chevron-right-double-c9H46Kl8.js} +1 -1
- zenml/zen_server/dashboard/assets/{cloud-only-BkUuI0lZ.js → cloud-only-BuP4Kt_7.js} +1 -1
- zenml/zen_server/dashboard/assets/code-browser-BJYErIjr.js +1 -0
- zenml/zen_server/dashboard/assets/codespaces-BitYDX9d.gif +0 -0
- zenml/zen_server/dashboard/assets/{copy-f3XGPPxt.js → copy-CaGlDsUy.js} +1 -1
- zenml/zen_server/dashboard/assets/create-stack-B2x2d4r1.js +1 -0
- zenml/zen_server/dashboard/assets/{docker-8uj__HHK.js → docker-BFAFXr2_.js} +1 -1
- zenml/zen_server/dashboard/assets/{dots-horizontal-sKQlWEni.js → dots-horizontal-C6K59vUm.js} +1 -1
- zenml/zen_server/dashboard/assets/flyte-Cj-xy_8I.svg +10 -0
- zenml/zen_server/dashboard/assets/form-schemas-Bap0f854.js +1 -0
- zenml/zen_server/dashboard/assets/gcp-Dj6ntk0L.js +1 -0
- zenml/zen_server/dashboard/assets/{help-FuHlZwn0.js → help-CwN931fX.js} +1 -1
- zenml/zen_server/dashboard/assets/{index-Bd1xgUQG.js → index-5GJ5ysEZ.js} +1 -1
- zenml/zen_server/dashboard/assets/{index-DaGknux4.css → index-6DYjZgDn.css} +1 -1
- zenml/zen_server/dashboard/assets/index-B9wVwe7u.js +55 -0
- zenml/zen_server/dashboard/assets/index-DFi8BroH.js +1 -0
- zenml/zen_server/dashboard/assets/{index.esm-DT4uyn2i.js → index.esm-BE1uqCX5.js} +1 -1
- zenml/zen_server/dashboard/assets/kubernetes-BjbR6D-1.js +1 -0
- zenml/zen_server/dashboard/assets/{layout-D6oiSbfd.js → layout-Dru15_XR.js} +1 -1
- zenml/zen_server/dashboard/assets/link-external-BT2L8hAQ.js +1 -0
- zenml/zen_server/dashboard/assets/{login-mutation-13A_JSVA.js → login-mutation-DwxUz8VA.js} +1 -1
- zenml/zen_server/dashboard/assets/{logs-CgeE2vZP.js → logs-GiDJXbLS.js} +1 -1
- zenml/zen_server/dashboard/assets/metaflow-weOkWNyT.svg +10 -0
- zenml/zen_server/dashboard/assets/{not-found-B0Mmb90p.js → not-found-D5i9DunU.js} +1 -1
- zenml/zen_server/dashboard/assets/{package-DdkziX79.js → package-DYKZ5jKW.js} +1 -1
- zenml/zen_server/dashboard/assets/page-BFuJICXM.js +9 -0
- zenml/zen_server/dashboard/assets/{page-BGwA9B1M.js → page-BiF8hLbO.js} +1 -1
- zenml/zen_server/dashboard/assets/{page-DugsjcQ_.js → page-BitfWsiW.js} +1 -1
- zenml/zen_server/dashboard/assets/page-CDOQLrPC.js +1 -0
- zenml/zen_server/dashboard/assets/page-CEJWu1YO.js +1 -0
- zenml/zen_server/dashboard/assets/page-CIbehp7V.js +1 -0
- zenml/zen_server/dashboard/assets/page-CLiRGfWo.js +1 -0
- zenml/zen_server/dashboard/assets/page-CV44mQn9.js +1 -0
- zenml/zen_server/dashboard/assets/page-CrSdkteO.js +2 -0
- zenml/zen_server/dashboard/assets/page-D5F3DJjm.js +1 -0
- zenml/zen_server/dashboard/assets/page-DE03uZZR.js +1 -0
- zenml/zen_server/dashboard/assets/page-DFCK65G9.js +1 -0
- zenml/zen_server/dashboard/assets/{page-RnG-qhv9.js → page-DGMa3ZQL.js} +1 -1
- zenml/zen_server/dashboard/assets/page-DI-qTWrm.js +1 -0
- zenml/zen_server/dashboard/assets/page-DQGCHKrQ.js +1 -0
- zenml/zen_server/dashboard/assets/{page-DSTQnBk-.js → page-DQdwZZ9x.js} +1 -1
- zenml/zen_server/dashboard/assets/page-DgM-N9RL.js +1 -0
- zenml/zen_server/dashboard/assets/page-Dt8VgzbE.js +1 -0
- zenml/zen_server/dashboard/assets/{page-DLpOnf7u.js → page-J0s8Sq3N.js} +1 -1
- zenml/zen_server/dashboard/assets/page-WCQ659by.js +1 -0
- zenml/zen_server/dashboard/assets/page-bimkItOg.js +1 -0
- zenml/zen_server/dashboard/assets/{page-hQaiQXfg.js → page-iwoJnwPv.js} +1 -1
- zenml/zen_server/dashboard/assets/{page-YiF_fNbe.js → page-oS4hqS8M.js} +1 -1
- zenml/zen_server/dashboard/assets/page-oSqx9dkH.js +1 -0
- zenml/zen_server/dashboard/assets/page-p3GqEAUW.js +1 -0
- zenml/zen_server/dashboard/assets/page-qvcUVPE-.js +1 -0
- zenml/zen_server/dashboard/assets/page-xQG6GmFJ.js +1 -0
- zenml/zen_server/dashboard/assets/{persist-3-5nOJ6m.js → persist-mEZN_fgH.js} +1 -1
- zenml/zen_server/dashboard/assets/persist-xsYgVtR1.js +1 -0
- zenml/zen_server/dashboard/assets/{plus-FB9-lEq_.js → plus-Bc8eLSDM.js} +1 -1
- zenml/zen_server/dashboard/assets/{refresh-COb6KYDi.js → refresh-hfgWPeto.js} +1 -1
- zenml/zen_server/dashboard/assets/rocket-SESCGQXm.js +1 -0
- zenml/zen_server/dashboard/assets/sharedSchema-BfZcy7aP.js +14 -0
- zenml/zen_server/dashboard/assets/stack-detail-query-CU4egfhp.js +1 -0
- zenml/zen_server/dashboard/assets/templates-1S_8WeSK.webp +0 -0
- zenml/zen_server/dashboard/assets/{trash-Cd5CSFqA.js → trash-DUWZWzse.js} +1 -1
- zenml/zen_server/dashboard/assets/{update-server-settings-mutation-B8GB_ubU.js → update-server-settings-mutation-DNqmQXDM.js} +1 -1
- zenml/zen_server/dashboard/assets/{url-hcMJkz8p.js → url-DwbuKk1b.js} +1 -1
- zenml/zen_server/dashboard/assets/{zod-CnykDKJj.js → zod-uFd1wBcd.js} +1 -1
- zenml/zen_server/dashboard/index.html +7 -7
- zenml/zen_server/dashboard_legacy/asset-manifest.json +4 -4
- zenml/zen_server/dashboard_legacy/index.html +1 -1
- zenml/zen_server/dashboard_legacy/{precache-manifest.9c473c96a43298343a7ce1256183123b.js → precache-manifest.290b95d5b43efa3368b3dc63d20c4782.js} +4 -4
- zenml/zen_server/dashboard_legacy/service-worker.js +1 -1
- zenml/zen_server/dashboard_legacy/static/js/{main.463c90b9.chunk.js → main.840d1bf0.chunk.js} +2 -2
- zenml/zen_server/dashboard_legacy/static/js/{main.463c90b9.chunk.js.map → main.840d1bf0.chunk.js.map} +1 -1
- zenml/zen_server/deploy/helm/Chart.yaml +1 -1
- zenml/zen_server/deploy/helm/README.md +2 -2
- zenml/zen_server/routers/service_connectors_endpoints.py +2 -4
- zenml/zen_server/routers/workspaces_endpoints.py +20 -66
- zenml/zen_server/secure_headers.py +120 -0
- zenml/zen_server/template_execution/runner_entrypoint_configuration.py +0 -2
- zenml/zen_server/template_execution/utils.py +1 -0
- zenml/zen_server/utils.py +0 -100
- zenml/zen_server/zen_server_api.py +4 -2
- zenml/zen_stores/migrations/versions/0.65.0_release.py +23 -0
- zenml/zen_stores/migrations/versions/0.66.0_release.py +23 -0
- zenml/zen_stores/migrations/versions/bf2120261b5a_add_configured_model_version_id.py +74 -0
- zenml/zen_stores/rest_zen_store.py +4 -21
- zenml/zen_stores/schemas/constants.py +16 -0
- zenml/zen_stores/schemas/model_schemas.py +9 -3
- zenml/zen_stores/schemas/pipeline_run_schemas.py +22 -8
- zenml/zen_stores/schemas/step_run_schemas.py +23 -12
- zenml/zen_stores/sql_zen_store.py +312 -300
- zenml/zen_stores/zen_store_interface.py +0 -16
- {zenml_nightly-0.64.0.dev20240811.dist-info → zenml_nightly-0.66.0.dev20240910.dist-info}/METADATA +10 -12
- {zenml_nightly-0.64.0.dev20240811.dist-info → zenml_nightly-0.66.0.dev20240910.dist-info}/RECORD +249 -217
- zenml/models/v2/misc/full_stack.py +0 -129
- zenml/new/pipelines/model_utils.py +0 -72
- zenml/zen_server/dashboard/assets/AlertDialogDropdownItem-ErO9aOgK.js +0 -1
- zenml/zen_server/dashboard/assets/AwarenessChannel-CLXo5rKM.js +0 -1
- zenml/zen_server/dashboard/assets/CollapsibleCard-BaUPiVg0.js +0 -1
- zenml/zen_server/dashboard/assets/Commands-JrcZK-3j.js +0 -1
- zenml/zen_server/dashboard/assets/CopyButton-Dbo52T1K.js +0 -2
- zenml/zen_server/dashboard/assets/EditSecretDialog-Bd7mFLS4.js +0 -1
- zenml/zen_server/dashboard/assets/ExecutionStatus-jH4OrWBq.js +0 -1
- zenml/zen_server/dashboard/assets/Infobox-BQ0aty32.js +0 -1
- zenml/zen_server/dashboard/assets/ProviderRadio-BBqkIuTd.js +0 -1
- zenml/zen_server/dashboard/assets/RadioItem-xLhXoiFV.js +0 -1
- zenml/zen_server/dashboard/assets/SearchField-C9R0mdaX.js +0 -1
- zenml/zen_server/dashboard/assets/SetPassword-52sNxNiO.js +0 -1
- zenml/zen_server/dashboard/assets/SuccessStep-DlkItqYG.js +0 -1
- zenml/zen_server/dashboard/assets/aws-0_3UsPif.js +0 -1
- zenml/zen_server/dashboard/assets/database-cXYNX9tt.js +0 -1
- zenml/zen_server/dashboard/assets/file-text-B9JibxTs.js +0 -1
- zenml/zen_server/dashboard/assets/index-DhIZtpxB.js +0 -55
- zenml/zen_server/dashboard/assets/page-7-v2OBm-.js +0 -1
- zenml/zen_server/dashboard/assets/page-B3ozwdD1.js +0 -1
- zenml/zen_server/dashboard/assets/page-BkjAUyTA.js +0 -1
- zenml/zen_server/dashboard/assets/page-BnacgBiy.js +0 -1
- zenml/zen_server/dashboard/assets/page-BxF_KMQ3.js +0 -2
- zenml/zen_server/dashboard/assets/page-C4POHC0K.js +0 -1
- zenml/zen_server/dashboard/assets/page-C9kudd44.js +0 -9
- zenml/zen_server/dashboard/assets/page-CA1j3GpJ.js +0 -1
- zenml/zen_server/dashboard/assets/page-CCY6yfmu.js +0 -1
- zenml/zen_server/dashboard/assets/page-CgTe7Bme.js +0 -1
- zenml/zen_server/dashboard/assets/page-Cgn-6v2Y.js +0 -1
- zenml/zen_server/dashboard/assets/page-CxQmQqDw.js +0 -1
- zenml/zen_server/dashboard/assets/page-D2Goey3H.js +0 -1
- zenml/zen_server/dashboard/assets/page-DTysUGOy.js +0 -1
- zenml/zen_server/dashboard/assets/page-D_EXUFJb.js +0 -1
- zenml/zen_server/dashboard/assets/page-Db15QzsM.js +0 -1
- zenml/zen_server/dashboard/assets/page-OFKSPyN7.js +0 -1
- zenml/zen_server/dashboard/assets/page-T2BtjwPl.js +0 -1
- zenml/zen_server/dashboard/assets/page-TXe1Eo3Z.js +0 -1
- zenml/zen_server/dashboard/assets/play-circle-XSkLR12B.js +0 -1
- zenml/zen_server/dashboard/assets/sharedSchema-BoYx_B_L.js +0 -14
- zenml/zen_server/dashboard/assets/stack-detail-query-B-US_-wa.js +0 -1
- zenml/zen_server/dashboard/assets/terminal-grtjrIEJ.js +0 -1
- {zenml_nightly-0.64.0.dev20240811.dist-info → zenml_nightly-0.66.0.dev20240910.dist-info}/LICENSE +0 -0
- {zenml_nightly-0.64.0.dev20240811.dist-info → zenml_nightly-0.66.0.dev20240910.dist-info}/WHEEL +0 -0
- {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}")
|