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.
@@ -371,6 +371,38 @@ 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
+```
+
+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,22 @@ def _configure_local_client(
                     "aws_session_token"
                 ] = credentials.token
 
+            aws_credentials_path = os.path.join(
+                users_home, ".aws", "credentials"
+            )
+
+            # Ensure the .aws directory exists
+            os.makedirs(os.path.dirname(aws_credentials_path), exist_ok=True)
+
+            # Ensure the credentials file exists
+            if not os.path.isfile(aws_credentials_path):
+                with open(aws_credentials_path, "a"):
+                    pass  # The file is automatically closed when exiting the with block
+
+            # Set the appropriate permissions for the .aws directory and credentials file
+            os.chmod(os.path.dirname(aws_credentials_path), 0o700)
+            os.chmod(aws_credentials_path, 0o600)
+
             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,65 @@ 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,
+                            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 +1117,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 +1274,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

@@ -36,6 +36,7 @@ class SkypilotAWSIntegration(Integration):
 
     NAME = SKYPILOT_AWS
     REQUIREMENTS = ["skypilot[aws]"]
+    APT_PACKAGES = ["openssh-client","rsync"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:
@@ -56,6 +57,7 @@ class SkypilotGCPIntegration(Integration):
 
     NAME = SKYPILOT_GCP
     REQUIREMENTS = ["skypilot[gcp]"]
+    APT_PACKAGES = ["openssh-client","rsync"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:
@@ -76,6 +78,7 @@ class SkypilotAzureIntegration(Integration):
 
     NAME = SKYPILOT_AZURE
     REQUIREMENTS = ["skypilot[azure]"]
+    APT_PACKAGES = ["openssh-client","rsync"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]: