Allow multi step configuration for skypilot

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

@@ -65,16 +65,16 @@ To use the SkyPilot VM Orchestrator, you need:
 * One of the SkyPilot integrations installed. You can install the SkyPilot integration for your cloud provider of choice using the following command:
   ```shell
     # For AWS
-    pip install "zenml[connectors-gcp]"
-    zenml integration install aws vm_aws 
+    pip install "zenml[connectors-aws]"
+    zenml integration install aws skypilot_aws 
 
     # for GCP
     pip install "zenml[connectors-gcp]"
-    zenml integration install gcp vm_gcp # for GCP
+    zenml integration install gcp skypilot_gcp # for GCP
 
     # for Azure
     pip install "zenml[connectors-azure]"
-    zenml integration install azure vm_azure # for Azure
+    zenml integration install azure skypilot_azure # for Azure
   ```
 * [Docker](https://www.docker.com) installed and running.
 * A [remote artifact store](../artifact-stores/artifact-stores.md) as part of your stack.
@@ -91,7 +91,7 @@ We need first to install the SkyPilot integration for AWS and the AWS connectors
 
   ```shell
     pip install "zenml[connectors-aws]"
-    zenml integration install aws vm_aws 
+    zenml integration install aws skypilot_aws 
   ```
 
 To provision VMs on AWS, your VM Orchestrator stack component needs to be configured to authenticate with [AWS Service Connector](../../../stacks-and-components/auth-management/aws-service-connector.md).
@@ -141,7 +141,7 @@ We need first to install the SkyPilot integration for GCP and the GCP extra for
 
   ```shell
     pip install "zenml[connectors-gcp]"
-    zenml integration install gcp vm_gcp
+    zenml integration install gcp skypilot_gcp
   ```
 
 To provision VMs on GCP, your VM Orchestrator stack component needs to be configured to authenticate with [GCP Service Connector](../../../stacks-and-components/auth-management/gcp-service-connector.md)
@@ -199,7 +199,7 @@ We need first to install the SkyPilot integration for Azure and the Azure extra
 
   ```shell
     pip install "zenml[connectors-azure]"
-    zenml integration install azure vm_azure 
+    zenml integration install azure skypilot_azure 
   ```
 
 To provision VMs on Azure, your VM Orchestrator stack component needs to be configured to authenticate with [Azure Service Connector](../../../stacks-and-components/auth-management/azure-service-connector.md)
@@ -371,6 +371,46 @@ skypilot_settings = SkypilotAzureOrchestratorSettings(
 {% endtab %}
 {% endtabs %}
 
+
+One of the key features of the SkyPilot VM Orchestrator is the ability to run each step of a pipeline on a separate VM with its own specific settings. This allows for fine-grained control over the resources allocated to each step, ensuring that each part of your pipeline has the necessary compute power while optimizing for cost and efficiency.
+
+## Configuring Step-Specific Resources
+
+The SkyPilot VM Orchestrator allows you to configure resources for each step individually. This means you can specify different VM types, CPU and memory requirements, and even use spot instances for certain steps while using on-demand instances for others.
+
+To configure step-specific resources, you can pass a `SkypilotBaseOrchestratorSettings` object to the `settings` parameter of the `@step` decorator. This object allows you to define various attributes such as `instance_type`, `cpus`, `memory`, `use_spot`, `region`, and more.
+
+Here's an example of how to configure specific resources for a step for the AWS cloud:
+
+```python
+from zenml.integrations.skypilot.flavors.skypilot_orchestrator_aws_vm_flavor import SkypilotAWSOrchestratorSettings
+
+# Settings for a specific step that requires more resources
+high_resource_settings = SkypilotAWSOrchestratorSettings(
+    instance_type='t2.2xlarge',
+    cpus=8,
+    memory=32,
+    use_spot=False,
+    region='us-east-1',
+    # ... other settings
+)
+
+@step(settings={"orchestrator.vm_aws": high_resource_settings})
+def my_resource_intensive_step():
+    # Step implementation
+    pass
+```
+
+{% hint style="warning" %}
+When configuring pipeline or step-specific resources, you can use the `settings`
+parameter to specifically target the orchestrator flavor you want to use
+`orchestrator.STACK_COMPONENT_FLAVOR` and not orchestrator component name
+`orchestrator.STACK_COMPONENT_NAME`. For example, if you want to configure
+resources for the `vm_gcp` flavor, you can use `settings={"orchestrator.vm_gcp": ...}`.
+{% endhint %}
+
+By using the `settings` parameter, you can tailor the resources for each step according to its specific needs. This flexibility allows you to optimize your pipeline execution for both performance and cost.
+
 Check out
 the [SDK docs](https://sdkdocs.zenml.io/latest/integration\_code\_docs/integrations-skypilot/#zenml.integrations.skypilot.flavors.skypilot\_orchestrator\_base\_vm\_flavor.SkypilotBaseOrchestratorSettings)
 for a full list of available attributes and [this docs page](/docs/book/user-guide/advanced-guide/pipelining-features/configure-steps-pipelines.md) for more

src/zenml/integrations/aws/service_connectors/aws_service_connector.py

@@ -27,6 +27,7 @@
 import base64
 import datetime
 import json
+import os
 import re
 from typing import Any, Dict, List, Optional, Tuple, cast
 
@@ -1371,6 +1372,21 @@ def _configure_local_client(
                     "aws_session_token"
                 ] = credentials.token
 
+            aws_credentials_path = os.path.join(
+                users_home, ".aws", "credentials"
+            )
+
+            # Create the file as well as the parent dir if needed.
+            dirname = os.path.split(aws_credentials_path)[0]
+            if not os.path.isdir(dirname):
+                os.makedirs(dirname)
+            with os.fdopen(
+                os.open(aws_credentials_path, os.O_WRONLY | os.O_CREAT, 0o600),
+                "w",
+            ):
+                pass
+
+            # Write the credentials to the file
             common.rewrite_credentials_file(all_profiles, users_home)
 
             logger.info(

src/zenml/integrations/gcp/service_connectors/gcp_service_connector.py

@@ -23,6 +23,9 @@
 import json
 import os
 import re
+import shutil
+import subprocess
+import tempfile
 from typing import Any, Dict, List, Optional, Tuple
 
 import google.api_core.exceptions
@@ -1026,6 +1029,7 @@ def _configure_local_client(
             NotImplementedError: If the connector instance does not support
                 local configuration for the configured resource type or
                 authentication method.registry
+            AuthorizationException: If the local client configuration fails.
         """
         resource_type = self.resource_type
 
@@ -1034,16 +1038,8 @@ def _configure_local_client(
 
             # There is no way to configure the local gcloud CLI to use
             # temporary OAuth 2.0 tokens. However, we can configure it to use
-            # the user account credentials or service account credentials
-            if self.auth_method == GCPAuthenticationMethods.USER_ACCOUNT:
-                assert isinstance(self.config, GCPUserAccountConfig)
-                # Use the user account credentials JSON to configure the
-                # local gcloud CLI
-                gcloud_config_json = (
-                    self.config.user_account_json.get_secret_value()
-                )
-
-            elif self.auth_method == GCPAuthenticationMethods.SERVICE_ACCOUNT:
+            # the service account credentials
+            if self.auth_method == GCPAuthenticationMethods.SERVICE_ACCOUNT:
                 assert isinstance(self.config, GCPServiceAccountConfig)
                 # Use the service account credentials JSON to configure the
                 # local gcloud CLI
@@ -1054,7 +1050,66 @@ def _configure_local_client(
             if gcloud_config_json:
                 from google.auth import _cloud_sdk
 
-                # Dump the user account or service account credentials JSON to
+                if not shutil.which("gcloud"):
+                    raise AuthorizationException(
+                        "The local gcloud CLI is not installed. Please "
+                        "install the gcloud CLI to use this feature."
+                    )
+
+                # Write the credentials JSON to a temporary file
+                with tempfile.NamedTemporaryFile(
+                    mode="w", suffix=".json", delete=True
+                ) as f:
+                    f.write(gcloud_config_json)
+                    f.flush()
+                    adc_path = f.name
+
+                    try:
+                        # Run the gcloud CLI command to configure the local
+                        # gcloud CLI to use the credentials JSON
+                        subprocess.run(
+                            [
+                                "gcloud",
+                                "auth",
+                                "login",
+                                "--cred-file",
+                                adc_path,
+                            ],
+                            check=True,
+                            stderr=subprocess.STDOUT,
+                            encoding="utf-8",
+                            stdout=subprocess.PIPE,
+                        )
+                    except subprocess.CalledProcessError as e:
+                        raise AuthorizationException(
+                            f"Failed to configure the local gcloud CLI to use "
+                            f"the credentials JSON: {e}\n"
+                            f"{e.stdout.decode()}"
+                        )
+
+                try:
+                    # Run the gcloud CLI command to configure the local gcloud
+                    # CLI to use the credentials project ID
+                    subprocess.run(
+                        [
+                            "gcloud",
+                            "config",
+                            "set",
+                            "project",
+                            self.config.project_id,
+                        ],
+                        check=True,
+                        stderr=subprocess.STDOUT,
+                        stdout=subprocess.PIPE,
+                    )
+                except subprocess.CalledProcessError as e:
+                    raise AuthorizationException(
+                        f"Failed to configure the local gcloud CLI to use "
+                        f"the project ID: {e}\n"
+                        f"{e.stdout.decode()}"
+                    )
+
+                # Dump the service account credentials JSON to
                 # the local gcloud application default credentials file
                 adc_path = (
                     _cloud_sdk.get_application_default_credentials_path()
@@ -1063,16 +1118,15 @@ def _configure_local_client(
                     f.write(gcloud_config_json)
 
                 logger.info(
-                    f"Updated the local gcloud default application "
-                    f"credentials file at '{adc_path}'"
+                    "Updated the local gcloud CLI and application default "
+                    f"credentials file ({adc_path})."
                 )
 
                 return
 
             raise NotImplementedError(
                 f"Local gcloud client configuration for resource type "
                 f"{resource_type} is only supported if the "
-                f"'{GCPAuthenticationMethods.USER_ACCOUNT}' or "
                 f"'{GCPAuthenticationMethods.SERVICE_ACCOUNT}' authentication "
                 f"method is used and only if the generation of temporary OAuth "
                 f"2.0 tokens is disabled by setting the "
@@ -1221,13 +1275,27 @@ def _auto_configure(
                     "GOOGLE_APPLICATION_CREDENTIALS"
                 )
                 if service_account_json_file is None:
-                    # Shouldn't happen since google.auth.default() should
-                    # already have loaded the credentials from the environment
+                    # No explicit service account JSON file was specified in the
+                    # environment, meaning that the credentials were loaded from
+                    # the GCP application default credentials (ADC) file.
+                    from google.auth import _cloud_sdk
+
+                    # Use the location of the gcloud application default
+                    # credentials file
+                    service_account_json_file = (
+                        _cloud_sdk.get_application_default_credentials_path()
+                    )
+
+                if not service_account_json_file or not os.path.isfile(
+                    service_account_json_file
+                ):
                     raise AuthorizationException(
-                        "No GCP service account credentials found in the "
-                        "environment. Please set the "
-                        "GOOGLE_APPLICATION_CREDENTIALS environment variable "
-                        "to the path of the service account JSON file."
+                        "No GCP service account credentials were found in the "
+                        "environment or the application default credentials"
+                        "path. Please set the GOOGLE_APPLICATION_CREDENTIALS "
+                        "environment variable to the path of the service "
+                        "account JSON file or run 'gcloud auth application-"
+                        "default login' to generate a new ADC file."
                     )
                 with open(service_account_json_file, "r") as f:
                     service_account_json = f.read()

src/zenml/integrations/skypilot/__init__.py

@@ -30,52 +30,52 @@
 SKYPILOT_GCP_ORCHESTRATOR_FLAVOR = "vm_gcp"
 SKYPILOT_AZURE_ORCHESTRATOR_FLAVOR = "vm_azure"
 
+class SkypilotGCPIntegration(Integration):
+    """Definition of Skypilot Integration for ZenML."""
 
-class SkypilotAWSIntegration(Integration):
-    """Definition of Skypilot AWS Integration for ZenML."""
-
-    NAME = SKYPILOT_AWS
-    REQUIREMENTS = ["skypilot[aws]"]
+    NAME = SKYPILOT_GCP
+    REQUIREMENTS = ["skypilot[aws,gcp,azure]<=0.4.1"]
+    APT_PACKAGES = ["openssh-client","rsync"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:
-        """Declare the stack component flavors for the Skypilot AWS integration.
+        """Declare the stack component flavors for the Skypilot GCP integration.
 
         Returns:
             List of stack component flavors for this integration.
         """
         from zenml.integrations.skypilot.flavors import (
-            SkypilotAWSOrchestratorFlavor,
+            SkypilotGCPOrchestratorFlavor,
         )
 
-        return [SkypilotAWSOrchestratorFlavor]
-
+        return [SkypilotGCPOrchestratorFlavor]
 
-class SkypilotGCPIntegration(Integration):
-    """Definition of Skypilot Integration for ZenML."""
+class SkypilotAWSIntegration(Integration):
+    """Definition of Skypilot AWS Integration for ZenML."""
 
-    NAME = SKYPILOT_GCP
-    REQUIREMENTS = ["skypilot[gcp]"]
+    NAME = SKYPILOT_AWS
+    REQUIREMENTS = ["skypilot[aws,gcp,azure]<=0.4.1"]
+    APT_PACKAGES = ["openssh-client","rsync"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:
-        """Declare the stack component flavors for the Skypilot GCP integration.
+        """Declare the stack component flavors for the Skypilot AWS integration.
 
         Returns:
             List of stack component flavors for this integration.
         """
         from zenml.integrations.skypilot.flavors import (
-            SkypilotGCPOrchestratorFlavor,
+            SkypilotAWSOrchestratorFlavor,
         )
 
-        return [SkypilotGCPOrchestratorFlavor]
-
+        return [SkypilotAWSOrchestratorFlavor]
 
 class SkypilotAzureIntegration(Integration):
     """Definition of Skypilot Integration for ZenML."""
 
     NAME = SKYPILOT_AZURE
-    REQUIREMENTS = ["skypilot[azure]"]
+    REQUIREMENTS = ["skypilot[aws,gcp,azure]<=0.4.1"]
+    APT_PACKAGES = ["openssh-client","rsync"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

src/zenml/integrations/skypilot/flavors/skypilot_orchestrator_base_vm_config.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Skypilot orchestrator base config and settings."""
 
-from typing import Dict, Literal, Optional, Union
+from typing import Dict, List, Literal, Optional, Union
 
 from zenml.config.base_settings import BaseSettings
 from zenml.logger import get_logger
@@ -107,6 +107,8 @@ class SkypilotBaseOrchestratorSettings(BaseSettings):
     down: bool = True
     stream_logs: bool = True
 
+    docker_run_args: List[str] = []
+
 
 class SkypilotBaseOrchestratorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
     BaseOrchestratorConfig, SkypilotBaseOrchestratorSettings