Skypilot: Lambda Edition

docs/book/stacks-and-components/component-guide/orchestrators/skypilot-vm.md

@@ -50,6 +50,14 @@ Read more about how to configure step-specific resources [here](#configuring-ste
 The SkyPilot VM Orchestrator does not currently support the ability to [schedule pipelines runs](/docs/book/user-guide/advanced-guide/pipelining-features/schedule-pipeline-runs.md)
 {% endhint %}
 
+{% hint style="info" %}
+All ZenML pipeline runs are executed using Docker containers within the VMs provisioned by the orchestrator.
+For that reason, you may need to configure your pipeline settings with `docker_run_args=["--gpus=all"]`
+to enable GPU support in the Docker container. 
+{% endhint %}
+
+{% hint style="info" %}
+
 
 ## How to deploy it
 
@@ -240,6 +248,36 @@ zenml orchestrator connect <ORCHESTRATOR_NAME> --connector azure-skypilot-vm
 zenml stack register <STACK_NAME> -o <ORCHESTRATOR_NAME> ... --set
 ```
 {% endtab %}
+
+{% tab title="Lambda Labs" %}
+
+[Lambda Labs](https://lambdalabs.com/service/gpu-cloud) is a cloud provider that offers GPU instances for machine learning workloads. Unlike the major cloud providers, with Lambda Labs we don't need to configure a service connector to authenticate with the cloud provider. Instead, we can directly use API keys to authenticate with the Lambda Labs API.
+
+  ```shell
+    zenml integration install skypilot_lambda
+  ```
+
+Once the integration is installed, we can register the orchestrator with the following command:
+
+```shell
+# For more secure and recommended way, we will register the API key as a secret
+zenml secret create lambda_api_key --scope user --api_key=<VALUE_1>
+# Register the orchestrator
+zenml orchestrator register <ORCHESTRATOR_NAME> --flavor vm_lambda --api_key={{lambda_api_key.api_key}}
+# Register and activate a stack with the new orchestrator
+zenml stack register <STACK_NAME> -o <ORCHESTRATOR_NAME> ... --set
+```
+
+{% hint style="info" %}
+The Lambda Labs orchestrator does not support some of the features like `spot_recovery`, `disk_tier`, `image_id`, `zone`, `idle_minutes_to_autostop`, `disk_size`, `use_spot`. It is recommended not to use these features with the Lambda Labs orchestrator and not to use [step-specific settings](#configuring-step-specific-resources).
+{% endhint %}
+
+{% hint style="warning" %}
+While testing the orchestrator, we noticed that the Lambda Labs orchestrator does not support the `down` flag. This means the orchestrator will not automatically tear down the cluster after all jobs finish. We recommend manually tearing down the cluster after all jobs finish to avoid unnecessary costs.
+{% endhint %}
+
+{% endtab %}
+
 {% endtabs %}
 
 #### Additional Configuration
@@ -263,6 +301,7 @@ For additional configuration of the Skypilot orchestrator, you can pass `Setting
 * `idle_minutes_to_autostop`: Automatically stop the cluster after this many minutes of idleness, i.e., no running or pending jobs in the cluster's job queue. Idleness gets reset whenever setting-up/running/pending jobs are found in the job queue. Setting this flag is equivalent to running `sky.launch(..., detach_run=True, ...)` and then `sky.autostop(idle_minutes=<minutes>)`. If not set, the cluster will not be autostopped.
 * `down`: Tear down the cluster after all jobs finish (successfully or abnormally). If `idle_minutes_to_autostop` is also set, the cluster will be torn down after the specified idle time. Note that if errors occur during provisioning/data syncing/setting up, the cluster will not be torn down for debugging purposes.
 * `stream_logs`: If True, show the logs in the terminal as they are generated while the cluster is running.
+* `docker_run_args`: Additional arguments to pass to the `docker run` command. For example, `['--gpus=all']` to use all GPUs available on the VM.
 
 The following code snippets show how to configure the orchestrator settings for each cloud provider:
 
@@ -291,6 +330,7 @@ skypilot_settings = SkypilotAWSOrchestratorSettings(
     idle_minutes_to_autostop=60,
     down=True,
     stream_logs=True
+    docker_run_args=["--gpus=all"]
 )
 
 
@@ -374,6 +414,34 @@ skypilot_settings = SkypilotAzureOrchestratorSettings(
 )
 ```
 
+{% endtab %}
+
+{% tab title="Lambda" %}
+
+**Code Example:**
+
+```python
+from zenml.integrations.skypilot_lambda import SkypilotLambdaOrchestratorSettings
+
+
+skypilot_settings = SkypilotLambdaOrchestratorSettings(
+    instance_type="gpu_1x_h100_pcie",
+    cluster_name="my_cluster",
+    retry_until_up=True,
+    idle_minutes_to_autostop=60,
+    down=True,
+    stream_logs=True,
+    docker_run_args=["--gpus=all"]
+)
+
+
+@pipeline(
+    settings={
+        "orchestrator.vm_lambda": skypilot_settings
+    }
+)
+```
+
 {% endtab %}
 {% endtabs %}
 

src/zenml/integrations/__init__.py

@@ -62,6 +62,7 @@
 from zenml.integrations.skypilot_aws import SkypilotAWSIntegration  # noqa
 from zenml.integrations.skypilot_gcp import SkypilotGCPIntegration  # noqa
 from zenml.integrations.skypilot_azure import SkypilotAzureIntegration  # noqa
+from zenml.integrations.skypilot_lambda import SkypilotLambdaIntegration  # noqa
 from zenml.integrations.slack import SlackIntegration  # noqa
 from zenml.integrations.spark import SparkIntegration  # noqa
 from zenml.integrations.tekton import TektonIntegration  # noqa

src/zenml/integrations/constants.py

@@ -57,6 +57,7 @@
 SKYPILOT_AWS = "skypilot_aws"
 SKYPILOT_GCP = "skypilot_gcp"
 SKYPILOT_AZURE = "skypilot_azure"
+SKYPILOT_LAMBDA = "skypilot_lambda"
 SLACK = "slack"
 SPARK = "spark"
 TEKTON = "tekton"

src/zenml/integrations/integration.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """Base and meta classes for ZenML integrations."""
 
+import re
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, cast
 
 import pkg_resources
@@ -83,15 +84,19 @@ def check_installation(cls) -> bool:
                         try:
                             requirements = dist.requires(extras=[extra])  # type: ignore[arg-type]
                         except pkg_resources.UnknownExtra as e:
-                            logger.debug("Unknown extra: " + str(e))
+                            logger.debug(f"Unknown extra: {str(e)}")
                             return False
                         deps.extend(requirements)
                 else:
                     deps = dist.requires()
 
                 for ri in deps:
                     try:
-                        pkg_resources.get_distribution(ri)
+                        # Remove the "extra == ..." part from the requirement string
+                        cleaned_req = re.sub(
+                            r"; extra == \"\w+\"", "", str(ri)
+                        )
+                        pkg_resources.get_distribution(cleaned_req)
                     except pkg_resources.DistributionNotFound as e:
                         logger.debug(
                             f"Unable to find required dependency "

src/zenml/integrations/skypilot/orchestrators/skypilot_base_vm_orchestrator.py

@@ -268,7 +268,7 @@ def prepare_or_run_pipeline(
         if docker_creds := stack.container_registry.credentials:
             docker_username, docker_password = docker_creds
             setup = (
-                f"docker login --username $DOCKER_USERNAME --password "
+                f"sudo docker login --username $DOCKER_USERNAME --password "
                 f"$DOCKER_PASSWORD {stack.container_registry.config.uri}"
             )
             task_envs = {
@@ -286,10 +286,14 @@ def prepare_or_run_pipeline(
 
         try:
             task = sky.Task(
-                run=f"docker run --rm {custom_run_args}{docker_environment_str} {image} {entrypoint_str} {arguments_str}",
+                run=f"sudo docker run --rm {custom_run_args}{docker_environment_str} {image} {entrypoint_str} {arguments_str}",
                 setup=setup,
                 envs=task_envs,
             )
+            logger.debug(
+                f"Running run: sudo docker run --rm {custom_run_args}{docker_environment_str} {image} {entrypoint_str} {arguments_str}"
+            )
+            logger.debug(f"Running run: {setup}")
             task = task.set_resources(
                 sky.Resources(
                     cloud=self.cloud,

src/zenml/integrations/skypilot_aws/__init__.py

@@ -31,7 +31,7 @@ class SkypilotAWSIntegration(Integration):
     """Definition of Skypilot AWS Integration for ZenML."""
 
     NAME = SKYPILOT_AWS
-    REQUIREMENTS = ["skypilot[aws]<=0.4.1"]
+    REQUIREMENTS = ["skypilot[aws]<=0.5.0"]
     APT_PACKAGES = ["openssh-client","rsync"]
 
     @classmethod

src/zenml/integrations/skypilot_azure/__init__.py

@@ -33,7 +33,7 @@ class SkypilotAzureIntegration(Integration):
     """Definition of Skypilot (Azure) Integration for ZenML."""
 
     NAME = SKYPILOT_AZURE
-    REQUIREMENTS = ["skypilot[azure]<=0.4.1"]
+    REQUIREMENTS = ["skypilot[azure]<=0.5.0"]
     APT_PACKAGES = ["openssh-client","rsync"]
 
     @classmethod

src/zenml/integrations/skypilot_gcp/__init__.py

@@ -30,7 +30,7 @@ class SkypilotGCPIntegration(Integration):
     """Definition of Skypilot (GCP) Integration for ZenML."""
 
     NAME = SKYPILOT_GCP
-    REQUIREMENTS = ["skypilot[gcp]<=0.4.1"]
+    REQUIREMENTS = ["skypilot[gcp]<=0.5.0"]
     APT_PACKAGES = ["openssh-client","rsync"]
 
     @classmethod

src/zenml/integrations/skypilot_lambda/__init__.py

@@ -0,0 +1,50 @@
+#  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 Skypilot Lambda integration for ZenML.
+
+The Skypilot integration sub-module powers an alternative to the local
+orchestrator for a remote orchestration of ZenML pipelines on VMs.
+"""
+from typing import List, Type
+
+from zenml.integrations.constants import (
+    SKYPILOT_LAMBDA,
+)
+from zenml.integrations.integration import Integration
+from zenml.stack import Flavor
+
+SKYPILOT_LAMBDA_ORCHESTRATOR_FLAVOR = "vm_lambda"
+
+
+class SkypilotLambdaIntegration(Integration):
+    """Definition of Skypilot Lambda Integration for ZenML."""
+
+    NAME = SKYPILOT_LAMBDA
+    REQUIREMENTS = ["skypilot[lambda]<=0.5.0"]
+
+    @classmethod
+    def flavors(cls) -> List[Type[Flavor]]:
+        """Declare the stack component flavors for the Skypilot Lambda integration.
+
+        Returns:
+            List of stack component flavors for this integration.
+        """
+        from zenml.integrations.skypilot_lambda.flavors import (
+            SkypilotLambdaOrchestratorFlavor,
+        )
+
+        return [SkypilotLambdaOrchestratorFlavor]
+
+
+SkypilotLambdaIntegration.check_installation()

src/zenml/integrations/skypilot_lambda/flavors/__init__.py

@@ -0,0 +1,26 @@
+#  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.
+"""Skypilot integration flavor for Skypilot Lambda orchestrator."""
+
+from zenml.integrations.skypilot_lambda.flavors.skypilot_orchestrator_lambda_vm_flavor import (
+    SkypilotLambdaOrchestratorConfig,
+    SkypilotLambdaOrchestratorFlavor,
+    SkypilotLambdaOrchestratorSettings,
+)
+
+__all__ = [
+    "SkypilotLambdaOrchestratorConfig",
+    "SkypilotLambdaOrchestratorFlavor",
+    "SkypilotLambdaOrchestratorSettings",
+]