Use docker credentials in the skypilot orchestrator (#1983)

* Use docker credentials in the skypilot orchestrator

* Update docs and add missing AWS connector requirements to S3 integration

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

@@ -95,7 +95,10 @@ We need first to install the SkyPilot integration for AWS and the AWS connectors
   ```
 
 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).
-To configure the AWS Service Connector, you need to register a new service connector, but first let's check the available service connector types using the following command:
+To configure the AWS Service Connector, you need to register a new service connector configured with AWS credentials that have at least the minimum permissions required by SkyPilot as documented [here](https://skypilot.readthedocs.io/en/latest/cloud-setup/cloud-permissions/aws.html).
+
+
+First, check that the AWS service connector type is available using the following command:
 
 ```
 zenml service-connector list-types --type aws
@@ -111,10 +114,10 @@ zenml service-connector list-types --type aws
 ┗━━━━━━━━━━━━━━━━━━━━━━━┷━━━━━━━━┷━━━━━━━━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━━━┷━━━━━━━┷━━━━━━━━┛
 ```
 
-For this example we will configure a service connector using the `iam-role` auth method. But before we can do that, we recommend you to create a new AWS profile that will be used by the service connector. Once we have created the profile, we can register a new service connector using the following command:
+Next, configure a service connector using the CLI or the dashboard with the AWS credentials. For example, the following command uses the local AWS CLI credentials to auto-configure the service connector:
 
 ```shell
-AWS_PROFILE=connectors zenml service-connector register aws-skypilot-vm --type aws --region=us-east-1 --auto-configure
+zenml service-connector register aws-skypilot-vm --type aws --region=us-east-1 --auto-configure
 ```
 
 This will automatically configure the service connector with the appropriate credentials and permissions to
@@ -162,7 +165,7 @@ For this example we will configure a service connector using the `user-account`
 login to GCP using the following command:
 
 ```shell
- gcloud auth application-default login 
+gcloud auth application-default login 
 ```
 
 This will open a browser window and ask you to login to your GCP account. Once you have logged in, you can register a new service connector using the

src/zenml/container_registries/base_container_registry.py

@@ -105,6 +105,18 @@ def credentials(self) -> Optional[Tuple[str, str]]:
         if secret:
             return secret.username, secret.password
 
+        connector = self.get_connector()
+        if connector:
+            from zenml.service_connectors.docker_service_connector import (
+                DockerServiceConnector,
+            )
+
+            if isinstance(connector, DockerServiceConnector):
+                return (
+                    connector.config.username.get_secret_value(),
+                    connector.config.password.get_secret_value(),
+                )
+
         return None
 
     @property

src/zenml/integrations/s3/__init__.py

@@ -41,7 +41,12 @@ class S3Integration(Integration):
     #
     # The above command installs boto3==1.26.76, so we use the same version
     # here to avoid the dependency resolution overhead.
-    REQUIREMENTS = ["s3fs>2022.3.0,<=2023.4.0", "boto3<=1.26.76"]
+    REQUIREMENTS = [
+        "s3fs>2022.3.0,<=2023.4.0",
+        "boto3<=1.26.76",
+        # The following dependencies are only required for the AWS connector.
+        "aws-profile-manager",
+    ]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

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

@@ -29,7 +29,6 @@
 
 if TYPE_CHECKING:
     from zenml.config.base_settings import BaseSettings
-    from zenml.stack import Stack
 
 logger = get_logger(__name__)
 
@@ -53,19 +52,6 @@ def cloud(self) -> sky.clouds.Cloud:
         """
         return sky.clouds.AWS()
 
-    def get_setup(self, stack: Optional["Stack"]) -> Optional[str]:
-        """Run to set up the sky job.
-
-        Args:
-            stack: The stack to use.
-
-        Returns:
-            A `setup` string.
-        """
-        assert stack is not None and stack.container_registry is not None
-        assert hasattr(stack.container_registry, "_get_region")
-        return f"aws ecr get-login-password --region {stack.container_registry._get_region()} | docker login --username AWS --password-stdin {stack.container_registry.config.uri}"
-
     @property
     def config(self) -> SkypilotAWSOrchestratorConfig:
         """Returns the `SkypilotAWSOrchestratorConfig` config.

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

@@ -126,17 +126,6 @@ def setup_credentials(self) -> None:
         assert connector is not None
         connector.configure_local_client()
 
-    def get_setup(self, stack: Optional["Stack"]) -> Optional[str]:
-        """Run to set up the sky job.
-
-        Args:
-            stack: The stack to use.
-
-        Returns:
-            A `setup` string.
-        """
-        return None
-
     @abstractmethod
     def prepare_environment_variable(self, set: bool = True) -> None:
         """Set up Environment variables that are required for the orchestrator.
@@ -200,12 +189,34 @@ def prepare_or_run_pipeline(
         # Set up credentials
         self.setup_credentials()
 
+        # Guaranteed by stack validation
+        assert stack is not None and stack.container_registry is not None
+
+        docker_creds = stack.container_registry.credentials
+        if docker_creds:
+            docker_username, docker_password = docker_creds
+            setup = (
+                f"docker login --username $DOCKER_USERNAME --password "
+                f"$DOCKER_PASSWORD {stack.container_registry.config.uri}"
+            )
+            task_envs = {
+                "DOCKER_USERNAME": docker_username,
+                "DOCKER_PASSWORD": docker_password,
+            }
+        else:
+            setup = None
+            task_envs = None
+
         # Run the entire pipeline
 
+        # Set the service connector AWS profile ENV variable
+        self.prepare_environment_variable(set=True)
+
         try:
             task = sky.Task(
                 run=f"docker run --rm {docker_environment_str} {image} {entrypoint_str} {arguments_str}",
-                setup=self.get_setup(stack),
+                setup=setup,
+                envs=task_envs,
             )
             task = task.set_resources(
                 sky.Resources(
@@ -237,8 +248,6 @@ def prepare_or_run_pipeline(
                             f"Found existing cluster {cluster_name}. Reusing..."
                         )
 
-            # Set the service connector AWS profile ENV variable
-            self.prepare_environment_variable(set=True)
             # Launch the cluster
             sky.launch(
                 task,