Lightening studio orchestrator (#2931)

* add lightening studio orchestrator

* update

* fix installing write zenml branch

* update the orchestrator

* fix pip install

* fix req

* update mypy

* fix machine usage

* update setting

* update setting

* fix deploy

* fix get settings

* add uv

* add async

* updated code

* update env setting

* remove breakpoints

* fixes

* run using a job

* update

* fix skypilot change

* Auto-update of NLP template

* update lightning

* Auto-update of Starter template

* update code

* update lighning orchestrator

* remove old wheel code

* cleanup imports

* fix async mode

* update async mode

* update lightning

* format

* chore: Upload code to main studio in lightning orchestrator entrypoint

* chore: Update lightning orchestrator entrypoint to upload code to main studio

* update ligning

* update docs

* update async

* fix provisoning

* add a cat logs for async

* ready lightning studio

* update based on reviews

* update docs and toc and init

* spell check

* add screenshots

* Optimised images with calibre/image-actions

* Refactor code utils

* Linting

* More linting

* Add mocked lib for api docs

* Some small fixes

---------

Co-authored-by: GitHub Actions <[email protected]>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Michael Schuster <[email protected]>

Author: 4 people

docs/book/.gitbook/assets/lightning_studio_ui.png

@@ -0,0 +1,205 @@
+
+---
+description: Orchestrating your pipelines to run on Lightning AI.
+---
+
+
+# Lightning AI Orchestrator
+
+[Lightning AI Studio](https://lightning.ai/) is a platform that simplifies the development and deployment of AI applications. The Lightning AI orchestrator is an integration provided by ZenML that allows you to run your pipelines on Lightning AI's infrastructure, leveraging its scalable compute resources and managed environment.
+
+
+{% hint style="warning" %}
+This component is only meant to be used within the context of a [remote ZenML deployment scenario](../../getting-started/deploying-zenml/README.md). Usage with a local ZenML deployment may lead to unexpected behavior!
+{% endhint %}
+
+
+## When to use it
+
+* You are looking for a fast and easy way to run your pipelines on GPU instances.
+
+* you're already using Lightning AI for your machine learning projects.
+
+* you want to leverage Lightning AI's managed infrastructure for running your pipelines.
+
+* you're looking for a solution that simplifies the deployment and scaling of your ML workflows.
+
+* you want to take advantage of Lightning AI's optimizations for machine learning workloads.
+
+
+## How to deploy it
+
+To use the [Lightning AI Studio](https://lightning.ai/) orchestrator, you need to have a Lightning AI account and the necessary credentials. You don't need to deploy any additional infrastructure, as the orchestrator will use Lightning AI's managed resources.
+
+
+## How it works
+
+The Lightning AI orchestrator is a ZenML orchestrator that runs your pipelines on Lightning AI's infrastructure. When you run a pipeline with the Lightning AI orchestrator, ZenML will archive your current ZenML repository and upload it to the Lightning AI studio. Once the code is archived, using `lightning-sdk`, ZenML will create a new stduio in Lightning AI and upload the code to it. Then ZenML runs list of commands via `studio.run()` to prepare for the pipeline run (e.g. installing dependencies, setting up the environment). Finally, ZenML will run the pipeline on Lightning AI's infrastructure.
+
+* You can always use an already existing studio by specifying the `main_studio_name` in the `LightningOrchestratorSettings`.
+* The orchestartor supports a async mode, which means that the pipeline will be run in the background and you can check the status of the run in the ZenML Dashboard or the Lightning AI Studio.
+* You can specify a list of custom commands that will be executed before running the pipeline. This can be useful for installing dependencies or setting up the environment.
+* The orchestrator supports both CPU and GPU machine types. You can specify the machine type in the `LightningOrchestratorSettings`.
+
+## How to use it
+
+To use the Lightning AI orchestrator, you need:
+
+*   The ZenML `lightning` integration installed. If you haven't done so, run
+
+    ```shell
+    pip install lightning-sdk
+    ```
+
+* A [remote artifact store](../artifact-stores/artifact-stores.md) as part of your stack.
+
+
+* [Lightning AI credentials](#lightning-ai-credentials)
+
+
+### Lightning AI credentials
+
+* `LIGHTNING_USER_ID`: Your Lightning AI user ID
+
+* `LIGHTNING_API_KEY`: Your Lightning AI API key
+
+* `LIGHTNING_USERNAME`: Your Lightning AI username (optional)
+
+* `LIGHTNING_TEAMSPACE`: Your Lightning AI teamspace (optional)
+
+* `LIGHTNING_ORG`: Your Lightning AI organization (optional)
+
+
+Alternatively, you can configure these credentials when registering the orchestrator.
+
+
+We can then register the orchestrator and use it in our active stack:
+
+```shell
+zenml orchestrator register lightning_orchestrator \
+    --flavor=lightning \
+    --user_id=<YOUR_LIGHTNING_USER_ID> \
+    --api_key=<YOUR_LIGHTNING_API_KEY> \
+    --username=<YOUR_LIGHTNING_USERNAME> \
+    --teamspace=<YOUR_LIGHTNING_TEAMSPACE> \
+    --organization=<YOUR_LIGHTNING_ORGANIZATION>
+
+# Register and activate a stack with the new orchestrator
+zenml stack register lightning_stack -o lightning_orchestrator ... --set
+```
+
+You can also configure the orchestrator at pipeline level, using the `orchestrator` parameter.
+
+
+```python
+from zenml.integrations.lightning.flavors.lightning_orchestrator_flavor import LightningOrchestratorSettings
+
+
+lightning_settings = LightningOrchestratorSettings(
+    main_studio_name="my_studio",
+    machine_type="cpu",
+    async_mode=True,
+    custom_commands=["pip install -r requirements.txt", "do something else"]
+)
+
+@pipeline(
+    settings={
+        "orchestrator.lightning": lightning_settings
+    }
+)
+def my_pipeline():
+    ...
+```
+
+
+{% hint style="info" %}
+ZenML will archive the current zenml repository (The code within the path where you run `zenml init`) and upload it to the Lightning AI studio.
+For this reason you need make sure that you are running `zenml init` in the same directory where you are running your pipeline.
+{% endhint %}
+
+![Lightning AI studio VSCode](../../.gitbook/assets/lightning_studio_vscode.png)
+
+{% hint style="info" %}
+The `custom_commands` attribute allows you to specify a list of shell commands that will be executed before running the pipeline. This can be useful for installing dependencies or setting up the environment, The commands will be executed in the root directory of the uploaded and extracted ZenML repository.
+{% endhint %}
+
+
+You can now run any ZenML pipeline using the Lightning AI orchestrator:
+
+```shell
+python file_that_runs_a_zenml_pipeline.py
+```
+
+
+### Lightning AI UI
+
+Lightning AI provides its own UI where you can monitor and manage your running applications, including the pipelines orchestrated by ZenML.
+![Lightning AI Studio](../../.gitbook/assets/lightning_studio_ui.png)
+
+
+For any runs executed on Lightning AI, you can get the URL to the Lightning AI UI in Python using the following code snippet:
+
+```python
+from zenml.client import Client
+
+pipeline_run = Client().get_pipeline_run("<PIPELINE_RUN_NAME>")
+orchestrator_url = pipeline_run.run_metadata["orchestrator_url"].value
+```
+
+
+### Additional configuration
+
+For additional configuration of the Lightning AI orchestrator, you can pass `LightningOrchestratorSettings` which allows you to configure various aspects of the Lightning AI execution environment:
+
+```python
+from zenml.integrations.lightning.flavors.lightning_orchestrator_flavor import LightningOrchestratorSettings
+
+lightning_settings = LightningOrchestratorSettings(
+    main_studio_name="my_studio",
+    machine_type="cpu",
+    async_mode=True,
+    custom_commands=["pip install -r requirements.txt", "do something else"]
+)
+```
+
+
+
+
+These settings can then be specified on either pipeline-level or step-level:
+
+```python
+# Either specify on pipeline-level
+@pipeline(
+    settings={
+        "orchestrator.lightning": lightning_settings
+    }
+)
+def my_pipeline():
+    ...
+
+# OR specify settings on step-level
+@step(
+    settings={
+        "orchestrator.lightning": lightning_settings
+    }
+)
+def my_step():
+    ...
+```
+
+Check out the [SDK docs](https://sdkdocs.zenml.io/latest/integration_code_docs/integrations-lightning/#zenml.integrations.lightning.flavors.lightning_orchestrator_flavor.LightningOrchestratorSettings) for a full list of available attributes and [this docs page](../../how-to/use-configuration-files/runtime-configuration.md) for more information on how to specify settings.
+
+
+To use GPUs with the Lightning AI orchestrator, you need to specify a GPU-enabled machine type in your settings:
+
+```python
+lightning_settings = LightningOrchestratorSettings(
+    machine_type="gpu", # or 
+)
+```
+
+
+Make sure to check Lightning AI's documentation for the available GPU-enabled machine types and their specifications.
+
+
+<figure><img src="https://static.scarf.sh/a.png?x-pxid=f0b4f458-0a54-4fcd-aa95-d5ee424815bc" alt="ZenML Scarf"><figcaption></figcaption></figure>

docs/book/.gitbook/assets/lightning_studio_vscode.png

@@ -203,6 +203,7 @@
   * [Airflow Orchestrator](component-guide/orchestrators/airflow.md)
   * [Skypilot VM Orchestrator](component-guide/orchestrators/skypilot-vm.md)
   * [HyperAI Orchestrator](component-guide/orchestrators/hyperai.md)
+  * [Lightning AI Orchestrator](component-guide/orchestrators/lightning.md)
   * [Develop a custom orchestrator](component-guide/orchestrators/custom.md)
 * [🏪 Artifact Stores](component-guide/artifact-stores/artifact-stores.md)
   * [Local Artifact Store](component-guide/artifact-stores/local.md)

docs/book/component-guide/orchestrators/lightning.md

@@ -251,5 +251,6 @@
     "IPython.core",
     "IPython.core.display",
     "IPython.core.display_functions",
-    "ipywidgets"
+    "ipywidgets",
+    "lightning_sdk"
 ]

docs/book/toc.md

@@ -464,6 +464,7 @@ module = [
     "huggingface_hub.*",
     "label_studio_sdk.*",
     "argilla.*",
+    "lightning_sdk.*",
     "peewee.*",
     "prodigy.*",
     "prodigy.components.*",

docs/mocked_libs.json

@@ -172,6 +172,7 @@ def handle_int_env_var(var: str, default: int = 0) -> int:
 )
 ENV_ZENML_IGNORE_FAILURE_HOOK = "ZENML_IGNORE_FAILURE_HOOK"
 ENV_ZENML_CUSTOM_SOURCE_ROOT = "ZENML_CUSTOM_SOURCE_ROOT"
+ENV_ZENML_WHEEL_PACKAGE_NAME = "ZENML_WHEEL_PACKAGE_NAME"
 
 # ZenML Server environment variables
 ENV_ZENML_SERVER_PREFIX = "ZENML_SERVER_"

pyproject.toml

@@ -219,7 +219,7 @@ def download_code_if_necessary(
                 code_reference=code_reference
             )
         elif code_path := deployment.code_path:
-            self.download_code_from_artifact_store(code_path=code_path)
+            code_utils.download_code_from_artifact_store(code_path=code_path)
         else:
             raise RuntimeError(
                 "Code download required but no code reference or path provided."
@@ -261,31 +261,6 @@ def download_code_from_code_repository(
         sys.path.insert(0, download_dir)
         os.chdir(download_dir)
 
-    def download_code_from_artifact_store(self, code_path: str) -> None:
-        """Download code from the artifact store.
-
-        Args:
-            code_path: Path where the code is stored.
-        """
-        logger.info(
-            "Downloading code from artifact store path `%s`.", code_path
-        )
-
-        # Do not remove this line, we need to instantiate the artifact store to
-        # register the filesystem needed for the file download
-        _ = Client().active_stack.artifact_store
-
-        extract_dir = os.path.abspath("code")
-        os.makedirs(extract_dir)
-
-        code_utils.download_and_extract_code(
-            code_path=code_path, extract_dir=extract_dir
-        )
-
-        source_utils.set_custom_source_root(extract_dir)
-        sys.path.insert(0, extract_dir)
-        os.chdir(extract_dir)
-
     def _should_download_code(
         self,
         deployment: "PipelineDeploymentResponse",

src/zenml/constants.py

@@ -36,6 +36,7 @@
 from zenml.integrations.great_expectations import (  # noqa
     GreatExpectationsIntegration,
 )
+from zenml.integrations.lightning import LightningIntegration  # noqa
 from zenml.integrations.huggingface import HuggingfaceIntegration  # noqa
 from zenml.integrations.hyperai import HyperAIIntegration  # noqa
 from zenml.integrations.kaniko import KanikoIntegration  # noqa

src/zenml/entrypoints/base_entrypoint_configuration.py

@@ -73,3 +73,4 @@
 VERTEX = "vertex"
 XGBOOST = "xgboost"
 VAULT = "vault"
+LIGHTNING = "lightning"

src/zenml/integrations/__init__.py

@@ -0,0 +1,48 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Initialization of the Lightning integration for ZenML."""
+
+from typing import List, Type
+
+from zenml.integrations.constants import (
+    LIGHTNING,
+)
+from zenml.integrations.integration import Integration
+from zenml.stack import Flavor
+
+LIGHTNING_ORCHESTRATOR_FLAVOR = "lightning"
+
+
+class LightningIntegration(Integration):
+    """Definition of Lightning Integration for ZenML."""
+
+    NAME = LIGHTNING
+    REQUIREMENTS = ["lightning-sdk"]
+
+    @classmethod
+    def flavors(cls) -> List[Type[Flavor]]:
+        """Declare the stack component flavors for the Lightning integration.
+
+        Returns:
+            List of stack component flavors for this integration.
+        """
+        from zenml.integrations.lightning.flavors import (
+            LightningOrchestratorFlavor,
+        )
+
+        return [
+            LightningOrchestratorFlavor,
+        ]
+
+LightningIntegration.check_installation()

src/zenml/integrations/constants.py

@@ -0,0 +1,23 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Lightning integration flavors."""
+
+from zenml.integrations.lightning.flavors.lightning_orchestrator_flavor import (
+    LightningOrchestratorConfig,
+    LightningOrchestratorFlavor,
+)
+__all__ = [
+    "LightningOrchestratorFlavor",
+    "LightningOrchestratorConfig",
+]