Repository: skypilot-org/skypilot
Pull request: Add docker support for SkyPilot
<!-- Describe the changes in this PR -->

This PR enable user-provided docker image on VM.

TODO:
- Test SSHProxyCommand on real-world usecase

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [x] Any manual or new tests for this PR (please specify below)
     - `pytest tests/test_smoke.py::test_job_queue_with_docker` 
     - `pytest tests/test_smoke.py::test_docker_preinstalled_package` 
- [x] All smoke tests: `pytest tests/test_smoke.py` 
- [ ] Relevant individual smoke tests:
- [ ] Backward compatibility tests: `bash tests/backward_comaptibility_tests.sh`


Problem statement:
[UX] Container support for task
We have received multiple requests from our users for the native container support in the task yaml, so that they don't need to manually run the `docker pull` and `docker run` in the `setup` and `run` sections separately.

Hints from issue discussion:
A blocker for one of our potential user.
Does this mean they would like the `run` commands to be executed within the docker container as well?
Yes, I suppose both setup and run sections will be run in the container, if that is posible.

Commit message: remove error merging

Patch to review:
diff --git a/docs/source/reference/yaml-spec.rst b/docs/source/reference/yaml-spec.rst
index 66d264189fe..3ed0c671831 100644
--- a/docs/source/reference/yaml-spec.rst
+++ b/docs/source/reference/yaml-spec.rst
@@ -102,8 +102,15 @@ Available fields:
         tpu_vm: False  # False to use TPU nodes (the default); True to use TPU VMs.
 
       # Custom image id (optional, advanced). The image id used to boot the
-      # instances. Only supported for AWS and GCP. If not specified, SkyPilot
-      # will use the default debian-based image suitable for machine learning tasks.
+      # instances. Only supported for AWS and GCP (for non-docker image). If not
+      # specified, SkyPilot will use the default debian-based image suitable for
+      # machine learning tasks.
+      #
+      # Docker support
+      # You can specify docker image to use by setting the image_id to
+      # `docker:<image name>` for Azure, AWS and GCP. For example,
+      #   image_id: docker:ubuntu:latest
+      # Currently, only debian and ubuntu images are supported.
       #
       # AWS
       # To find AWS AMI ids: https://leaherb.com/how-to-find-an-aws-marketplace-ami-image-id
diff --git a/sky/authentication.py b/sky/authentication.py
index 4a6d96a9f72..deb92fdcba2 100644
--- a/sky/authentication.py
+++ b/sky/authentication.py
@@ -225,23 +225,29 @@ def setup_gcp_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
             check=True,
             shell=True,
             stdout=subprocess.DEVNULL)
-        # Enable ssh port for all the instances
-        enable_ssh_cmd = ('gcloud compute firewall-rules create '
-                          'allow-ssh-ingress-from-iap '
-                          '--direction=INGRESS '
-                          '--action=allow '
-                          '--rules=tcp:22 '
-                          '--source-ranges=0.0.0.0/0')
-        proc = subprocess.run(enable_ssh_cmd,
-                              check=False,
-                              shell=True,
-                              stdout=subprocess.DEVNULL,
-                              stderr=subprocess.PIPE)
-        if proc.returncode != 0 and 'already exists' not in proc.stderr.decode(
-                'utf-8'):
-            subprocess_utils.handle_returncode(proc.returncode, enable_ssh_cmd,
-                                               'Failed to enable ssh port.',
-                                               proc.stderr.decode('utf-8'))
+
+        def allow_ssh_port(port):
+            # Enable ssh port for all the instances
+            enable_ssh_cmd = ('gcloud compute firewall-rules create '
+                              'allow-ssh-ingress-from-iap '
+                              '--direction=INGRESS '
+                              '--action=allow '
+                              f'--rules=tcp:{port} '
+                              '--source-ranges=0.0.0.0/0')
+            proc = subprocess.run(enable_ssh_cmd,
+                                  check=False,
+                                  shell=True,
+                                  stdout=subprocess.DEVNULL,
+                                  stderr=subprocess.PIPE)
+            if proc.returncode != 0 and \
+                'already exists' not in proc.stderr.decode('utf-8'):
+                subprocess_utils.handle_returncode(
+                    proc.returncode, enable_ssh_cmd,
+                    'Failed to enable ssh port.', proc.stderr.decode('utf-8'))
+
+        allow_ssh_port('22')
+        # For docker host machine
+        allow_ssh_port('10022')
         return config
 
     # OS Login is not enabled for the project. Add the ssh key directly to the
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index 2ba88bc05bd..fbcb0ae4330 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -402,6 +402,8 @@ def _get_generated_config(cls, autogen_comment: str, host_name: str,
             """.rstrip())
         codegen = codegen + '\n'
         return codegen
+        # TODO(tian): add a host_name + '-host' option with Port=10022
+        # to enable ssh into host machine when we are running with docker
 
     @classmethod
     @timeline.FileLockEvent(ssh_conf_lock_path)
@@ -777,6 +779,7 @@ def write_cluster_config(
         cluster_name: str,
         local_wheel_path: pathlib.Path,
         wheel_hash: str,
+        docker_image: Optional[str],
         region: Optional[clouds.Region] = None,
         zones: Optional[List[clouds.Zone]] = None,
         dryrun: bool = False,
@@ -854,6 +857,10 @@ def write_cluster_config(
             assert region_name in ssh_proxy_command_config, (
                 region_name, ssh_proxy_command_config)
             ssh_proxy_command = ssh_proxy_command_config[region_name]
+    public_key_path = os.path.expanduser(auth.PUBLIC_SSH_KEY_PATH)
+    with open(public_key_path, 'r') as f:
+        public_key = f.read()
+
     logger.debug(f'Using ssh_proxy_command: {ssh_proxy_command!r}')
 
     # User-supplied instance tags.
@@ -871,6 +878,10 @@ def write_cluster_config(
         f'open(os.path.expanduser("{constants.SKY_REMOTE_RAY_PORT_FILE}"), "w"))\''
     )
 
+    # pylint: disable=import-outside-toplevel
+    from sky.backends.docker_utils import \
+        DEFAULT_DOCKER_CONTAINER_NAME
+
     # Use a tmp file path to avoid incomplete YAML file being re-used in the
     # future.
     tmp_yaml_path = yaml_path + '.tmp'
@@ -906,6 +917,14 @@ def write_cluster_config(
                 # User-supplied instance tags.
                 'instance_tags': instance_tags,
 
+                # For docker authentication
+                'public_key': public_key,
+
+                # Docker config
+                'use_docker': docker_image is not None,
+                'docker_image': docker_image,
+                'docker_container_name': DEFAULT_DOCKER_CONTAINER_NAME,
+
                 # Azure only:
                 'azure_subscription_id': azure_subscription_id,
                 'resource_group': f'{cluster_name}-{region_name}',
@@ -1092,12 +1111,11 @@ def wait_until_ray_cluster_ready(
     num_nodes: int,
     log_path: str,
     is_local_cloud: bool = False,
+    use_docker: bool = False,
     nodes_launching_progress_timeout: Optional[int] = None,
-) -> bool:
-    """Returns whether the entire ray cluster is ready."""
-    if num_nodes <= 1:
-        return True
-
+) -> Tuple[bool, Optional[str]]:
+    """Returns whether the entire ray cluster is ready, and docker username
+    if launched with docker."""
     # Manually fetching head ip instead of using `ray exec` to avoid the bug
     # that `ray exec` fails to connect to the head node after some workers
     # launched especially for Azure.
@@ -1106,9 +1124,18 @@ def wait_until_ray_cluster_ready(
             cluster_config_file, max_attempts=WAIT_HEAD_NODE_IP_MAX_ATTEMPTS)
     except RuntimeError as e:
         logger.error(e)
-        return False  # failed
+        return False, None  # failed
+
+    docker_user = None
+    if use_docker:
+        # pylint: disable=import-outside-toplevel
+        from sky.backends.docker_utils import docker_host_setup
+        docker_user = docker_host_setup(head_ip, cluster_config_file)
+
+    if num_nodes <= 1:
+        return True, docker_user
 
-    ssh_credentials = ssh_credential_from_yaml(cluster_config_file)
+    ssh_credentials = ssh_credential_from_yaml(cluster_config_file, docker_user)
     last_nodes_so_far = 0
     start = time.time()
     runner = command_runner.SSHCommandRunner(head_ip, **ssh_credentials)
@@ -1167,7 +1194,7 @@ def wait_until_ray_cluster_ready(
                     'Timed out: waited for more than '
                     f'{nodes_launching_progress_timeout} seconds for new '
                     'workers to be provisioned, but no progress.')
-                return False  # failed
+                return False, None  # failed
 
             if '(no pending nodes)' in output and '(no failures)' in output:
                 # Bug in ray autoscaler: e.g., on GCP, if requesting 2 nodes
@@ -1177,16 +1204,18 @@ def wait_until_ray_cluster_ready(
                 logger.error(
                     'Failed to launch multiple nodes on '
                     'GCP due to a nondeterministic bug in ray autoscaler.')
-                return False  # failed
+                return False, None  # failed
             time.sleep(10)
-    return True  # success
+    return True, docker_user  # success
 
 
-def ssh_credential_from_yaml(cluster_yaml: str) -> Dict[str, str]:
+def ssh_credential_from_yaml(cluster_yaml: str,
+                             docker_user: Optional[str] = None
+                            ) -> Dict[str, str]:
     """Returns ssh_user, ssh_private_key and ssh_control name."""
     config = common_utils.read_yaml(cluster_yaml)
     auth_section = config['auth']
-    ssh_user = auth_section['ssh_user'].strip()
+    ssh_user = docker_user if docker_user else auth_section['ssh_user'].strip()
     ssh_private_key = auth_section.get('ssh_private_key')
     ssh_control_name = config.get('cluster_name', '__default__')
     ssh_proxy_command = auth_section.get('ssh_proxy_command')
@@ -2010,7 +2039,8 @@ def _update_cluster_status_no_lock(
             raise exceptions.FetchIPError(
                 reason=exceptions.FetchIPError.Reason.HEAD)
         # Check if ray cluster status is healthy.
-        ssh_credentials = ssh_credential_from_yaml(handle.cluster_yaml)
+        ssh_credentials = ssh_credential_from_yaml(handle.cluster_yaml,
+                                                   handle.docker_user)
         runner = command_runner.SSHCommandRunner(external_ips[0],
                                                  **ssh_credentials)
         rc, output, _ = runner.run(RAY_STATUS_WITH_SKY_RAY_PORT_COMMAND,
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index 2073a2f453f..b03d0e8ce42 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -1301,6 +1301,7 @@ def _retry_zones(
         cluster_name: str,
         cloud_user_identity: Optional[List[str]],
         prev_cluster_status: Optional[global_user_state.ClusterStatus],
+        docker_image: Optional[str],
     ):
         """The provision retry loop."""
         style = colorama.Style
@@ -1365,6 +1366,7 @@ def _retry_zones(
                     cluster_name,
                     self._local_wheel_path,
                     self._wheel_hash,
+                    docker_image,
                     region=region,
                     zones=zones,
                     dryrun=dryrun,
@@ -1421,7 +1423,8 @@ def _retry_zones(
             }
             status, stdout, stderr, head_ip = self._gang_schedule_ray_up(
                 to_provision.cloud, cluster_config_file, handle, log_abs_path,
-                stream_logs, logging_info, to_provision.use_spot)
+                stream_logs, logging_info, to_provision.use_spot, docker_image
+                is not None)
 
             if status == self.GangSchedulingStatus.CLUSTER_READY:
                 if cluster_exists:
@@ -1568,10 +1571,9 @@ def _tpu_pod_setup(self, cluster_yaml: str,
 
     @timeline.event
     def _gang_schedule_ray_up(
-            self, to_provision_cloud: clouds.Cloud, cluster_config_file: str,
-            cluster_handle: 'backends.CloudVmRayResourceHandle',
-            log_abs_path: str, stream_logs: bool, logging_info: dict,
-            use_spot: bool
+        self, to_provision_cloud: clouds.Cloud, cluster_config_file: str,
+        cluster_handle: 'backends.CloudVmRayResourceHandle', log_abs_path: str,
+        stream_logs: bool, logging_info: dict, use_spot: bool, use_docker: bool
     ) -> Tuple[GangSchedulingStatus, str, str, Optional[str]]:
         """Provisions a cluster via 'ray up' and wait until fully provisioned.
 
@@ -1748,6 +1750,7 @@ def need_ray_up(
             head_ip = None
             if len(ip_list) == 1:
                 head_ip = ip_list[0]
+            # If single node, docker_user is set up after handle is created.
             return (self.GangSchedulingStatus.CLUSTER_READY, stdout, stderr,
                     head_ip)
 
@@ -1777,7 +1780,8 @@ def need_ray_up(
             log_path=log_abs_path,
             nodes_launching_progress_timeout=_NODES_LAUNCHING_PROGRESS_TIMEOUT[
                 type(to_provision_cloud)],
-            is_local_cloud=isinstance(to_provision_cloud, clouds.Local))
+            is_local_cloud=isinstance(to_provision_cloud, clouds.Local),
+            use_docker=use_docker)
         if cluster_ready:
             cluster_status = self.GangSchedulingStatus.CLUSTER_READY
             # ray up --no-restart again with upscaling_speed=0 after cluster is
@@ -1893,7 +1897,8 @@ def provision_with_retries(
                     stream_logs=stream_logs,
                     cluster_name=cluster_name,
                     cloud_user_identity=cloud_user,
-                    prev_cluster_status=prev_cluster_status)
+                    prev_cluster_status=prev_cluster_status,
+                    docker_image=task.docker_image)
                 if dryrun:
                     return
             except (exceptions.InvalidClusterNameError,
@@ -1989,6 +1994,7 @@ class CloudVmRayResourceHandle(backends.backend.ResourceHandle):
         for all nodes in a cluster. Filled in after successful task execution.
     - (optional) Launched num nodes
     - (optional) Launched resources
+    - (optional) Docker user name
     - (optional) If TPU(s) are managed, a path to a deletion script.
     """
     _VERSION = 3
@@ -2012,6 +2018,7 @@ def __init__(self,
         self.stable_internal_external_ips = stable_internal_external_ips
         self.launched_nodes = launched_nodes
         self.launched_resources = launched_resources
+        self.docker_user: Optional[str] = None
         self.tpu_create_script = tpu_create_script
         self.tpu_delete_script = tpu_delete_script
         self._maybe_make_local_handle()
@@ -2163,6 +2170,26 @@ def get_hourly_price(self) -> float:
                        self.launched_nodes)
         return hourly_cost
 
+    def docker_setup(self, cluster_config_file: str):
+        ip_list = self.external_ips()
+        assert ip_list is not None
+        # pylint: disable=import-outside-toplevel
+        from sky.backends.docker_utils import docker_host_setup
+        if len(ip_list) == 1:
+            # Setup head node.
+            docker_user = docker_host_setup(ip_list[0], cluster_config_file)
+            self.docker_user = docker_user
+        else:
+            # Setup docker for all worker nodes. Head node setup is done
+            # in `backend_utils.wait_until_ray_cluster_ready`.
+            # Skip head ip here.
+            for worker_ip in ip_list[1:]:
+                docker_user = docker_host_setup(worker_ip, cluster_config_file)
+                if self.docker_user is not None:
+                    assert docker_user == self.docker_user
+                else:
+                    self.docker_user = docker_user
+
     @property
     def cluster_yaml(self):
         return os.path.expanduser(self._cluster_yaml)
@@ -2446,6 +2473,9 @@ def _provision(
                                           use_cached_ips=False)
             assert ip_list is not None, handle
 
+            if task.docker_image is not None:
+                handle.docker_setup(cluster_config_file)
+
             if 'tpu_name' in config_dict:
                 self._set_tpu_name(handle, config_dict['tpu_name'])
 
@@ -2463,7 +2493,7 @@ def _provision(
                     handle.launched_resources.cloud.get_zone_shell_cmd())
                 if get_zone_cmd is not None:
                     ssh_credentials = backend_utils.ssh_credential_from_yaml(
-                        handle.cluster_yaml)
+                        handle.cluster_yaml, handle.docker_user)
                     runners = command_runner.SSHCommandRunner.make_runner_list(
                         ip_list, **ssh_credentials)
 
@@ -2528,7 +2558,8 @@ def _update_after_cluster_provisioned(
             # update_status will query the ray job status for all INIT /
             # PENDING / RUNNING jobs for the real status, since we do not
             # know the actual previous status of the cluster.
-            job_owner = onprem_utils.get_job_owner(handle.cluster_yaml)
+            job_owner = onprem_utils.get_job_owner(handle.cluster_yaml,
+                                                   handle.docker_user)
             cmd = job_lib.JobLibCodeGen.update_status(job_owner)
             with log_utils.safe_rich_status('[bold cyan]Preparing Job Queue'):
                 returncode, _, stderr = self.run_on_head(handle,
@@ -2563,6 +2594,8 @@ def _update_after_cluster_provisioned(
             usage_lib.messages.usage.update_final_cluster_status(
                 global_user_state.ClusterStatus.UP)
             auth_config = common_utils.read_yaml(handle.cluster_yaml)['auth']
+            if handle.docker_user:
+                auth_config['ssh_user'] = handle.docker_user
             backend_utils.SSHConfigHelper.add_cluster(handle.cluster_name,
                                                       ip_list, auth_config)
 
@@ -2601,7 +2634,7 @@ def _sync_workdir(self, handle: CloudVmRayResourceHandle,
         log_path = os.path.join(self.log_dir, 'workdir_sync.log')
 
         ssh_credentials = backend_utils.ssh_credential_from_yaml(
-            handle.cluster_yaml)
+            handle.cluster_yaml, handle.docker_user)
 
         # TODO(zhwu): refactor this with backend_utils.parallel_cmd_with_rsync
         runners = command_runner.SSHCommandRunner.make_runner_list(
@@ -2661,7 +2694,7 @@ def _setup(self, handle: CloudVmRayResourceHandle, task: task_lib.Task,
             ip_list = handle.external_ips()
             assert ip_list is not None, 'external_ips is not cached in handle'
             ssh_credentials = backend_utils.ssh_credential_from_yaml(
-                handle.cluster_yaml)
+                handle.cluster_yaml, handle.docker_user)
             # Disable connection sharing for setup script to avoid old
             # connections being reused, which may cause stale ssh agent
             # forwarding.
@@ -2747,7 +2780,7 @@ def _exec_code_on_head(
         style = colorama.Style
         fore = colorama.Fore
         ssh_credentials = backend_utils.ssh_credential_from_yaml(
-            handle.cluster_yaml)
+            handle.cluster_yaml, handle.docker_user)
         runner = command_runner.SSHCommandRunner(handle.head_ip,
                                                  **ssh_credentials)
         with tempfile.NamedTemporaryFile('w', prefix='sky_app_') as fp:
@@ -2861,7 +2894,7 @@ def _setup_and_create_job_cmd_on_local_head(
     ):
         """Generates and prepares job submission code for local clusters."""
         ssh_credentials = backend_utils.ssh_credential_from_yaml(
-            handle.cluster_yaml)
+            handle.cluster_yaml, handle.docker_user)
         ssh_user = ssh_credentials['ssh_user']
         runner = command_runner.SSHCommandRunner(handle.head_ip,
                                                  **ssh_credentials)
@@ -3066,7 +3099,8 @@ def get_job_status(
 
     def cancel_jobs(self, handle: CloudVmRayResourceHandle,
                     jobs: Optional[List[int]]):
-        job_owner = onprem_utils.get_job_owner(handle.cluster_yaml)
+        job_owner = onprem_utils.get_job_owner(handle.cluster_yaml,
+                                               handle.docker_user)
         code = job_lib.JobLibCodeGen.cancel_jobs(job_owner, jobs)
 
         # All error messages should have been redirected to stdout.
@@ -3124,7 +3158,7 @@ def sync_down_logs(
         ip_list = handle.external_ips()
         assert ip_list is not None, 'external_ips is not cached in handle'
         ssh_credentials = backend_utils.ssh_credential_from_yaml(
-            handle.cluster_yaml)
+            handle.cluster_yaml, handle.docker_user)
         runners = command_runner.SSHCommandRunner.make_runner_list(
             ip_list, **ssh_credentials)
 
@@ -3162,7 +3196,8 @@ def tail_logs(self,
                   job_id: Optional[int],
                   spot_job_id: Optional[int] = None,
                   follow: bool = True) -> int:
-        job_owner = onprem_utils.get_job_owner(handle.cluster_yaml)
+        job_owner = onprem_utils.get_job_owner(handle.cluster_yaml,
+                                               handle.docker_user)
         code = job_lib.JobLibCodeGen.tail_logs(job_owner,
                                                job_id,
                                                spot_job_id=spot_job_id,
@@ -3499,6 +3534,8 @@ def post_teardown_cleanup(self,
         # be removed after the cluster entry in the database is removed.
         config = common_utils.read_yaml(handle.cluster_yaml)
         auth_config = config['auth']
+        if handle.docker_user:
+            auth_config['ssh_user'] = handle.docker_user
         backend_utils.SSHConfigHelper.remove_cluster(handle.cluster_name,
                                                      handle.head_ip,
                                                      auth_config)
@@ -3587,7 +3624,7 @@ def run_on_head(
         head_ip = backend_utils.get_head_ip(handle, use_cached_head_ip,
                                             _FETCH_IP_MAX_ATTEMPTS)
         ssh_credentials = backend_utils.ssh_credential_from_yaml(
-            handle.cluster_yaml)
+            handle.cluster_yaml, handle.docker_user)
         runner = command_runner.SSHCommandRunner(head_ip, **ssh_credentials)
         if under_remote_workdir:
             cmd = f'cd {SKY_REMOTE_WORKDIR} && {cmd}'
@@ -3669,7 +3706,7 @@ def _set_tpu_name(self, handle: CloudVmRayResourceHandle,
         ip_list = handle.external_ips()
         assert ip_list is not None, 'external_ips is not cached in handle'
         ssh_credentials = backend_utils.ssh_credential_from_yaml(
-            handle.cluster_yaml)
+            handle.cluster_yaml, handle.docker_user)
 
         runners = command_runner.SSHCommandRunner.make_runner_list(
             ip_list, **ssh_credentials)
@@ -3705,7 +3742,7 @@ def _execute_file_mounts(self, handle: CloudVmRayResourceHandle,
         ip_list = handle.external_ips()
         assert ip_list is not None, 'external_ips is not cached in handle'
         ssh_credentials = backend_utils.ssh_credential_from_yaml(
-            handle.cluster_yaml)
+            handle.cluster_yaml, handle.docker_user)
         runners = command_runner.SSHCommandRunner.make_runner_list(
             ip_list, **ssh_credentials)
         log_path = os.path.join(self.log_dir, 'file_mounts.log')
@@ -3852,7 +3889,7 @@ def _execute_storage_mounts(self, handle: CloudVmRayResourceHandle,
         ip_list = handle.external_ips()
         assert ip_list is not None, 'external_ips is not cached in handle'
         ssh_credentials = backend_utils.ssh_credential_from_yaml(
-            handle.cluster_yaml)
+            handle.cluster_yaml, handle.docker_user)
         runners = command_runner.SSHCommandRunner.make_runner_list(
             ip_list, **ssh_credentials)
         log_path = os.path.join(self.log_dir, 'storage_mounts.log')
diff --git a/sky/backends/docker_utils.py b/sky/backends/docker_utils.py
index 58a10f6c34c..1c694d689c6 100644
--- a/sky/backends/docker_utils.py
+++ b/sky/backends/docker_utils.py
@@ -1,10 +1,11 @@
 """Utilities for docker image generation."""
+import json
 import os
 import shutil
 import subprocess
 import tempfile
 import textwrap
-from typing import Dict, Optional, Tuple
+from typing import Dict, Optional, Tuple, Any
 
 import colorama
 
@@ -12,6 +13,10 @@
 from sky import sky_logging
 from sky import task as task_mod
 
+from ray.autoscaler._private.cli_logger import cli_logger
+from ray.autoscaler._private.command_runner import DockerCommandRunner
+from ray.autoscaler._private.docker import docker_start_cmds
+
 logger = sky_logging.init_logger(__name__)
 
 # Add docker-cli from official docker image to support docker-in-docker.
@@ -27,6 +32,10 @@
 DOCKERFILE_COPYCMD = """COPY {copy_command}"""
 DOCKERFILE_RUNCMD = """CMD {run_command}"""
 
+# Docker default options
+DEFAULT_DOCKER_IMAGE = 'brymastr/ubuntu-rsync'
+DEFAULT_DOCKER_CONTAINER_NAME = 'sky_container'
+
 CONDA_SETUP_PREFIX = '. $(conda info --base)/etc/profile.d/conda.sh 2> ' \
                      '/dev/null || true '
 
@@ -144,6 +153,63 @@ def _execute_build(tag, context_path):
         raise
 
 
+def extract_docker_image_from_resources(
+        resources: Optional[Dict[str, Any]]) -> Optional[str]:
+    """
+    Extracts the docker image from the resources dict.
+    """
+    if resources is None:
+        return None
+    if 'image_id' not in resources:
+        return None
+    if isinstance(resources['image_id'], str):
+        if resources['image_id'].startswith('docker:'):
+            image = resources['image_id'][len('docker:'):]
+            del resources['image_id']
+            return image
+    return None
+
+
+def docker_host_setup(ip: str, cluster_config_file: str) -> str:
+    """
+    Performs setup on the docker host. This setup do 2 things:
+      1. Find docker container username.
+      2. Change the host ssh port to 10022 and enable ssh login
+         to the docker container.
+    """
+    # pylint: disable=import-outside-toplevel
+    from sky.backends.backend_utils import ssh_credential_from_yaml
+    from sky.utils import command_runner
+    ssh_credentials = ssh_credential_from_yaml(cluster_config_file)
+    runner = command_runner.SSHCommandRunner(ip, **ssh_credentials)
+    whoami_returncode, whoami_stdout, whoami_stderr = runner.run(
+        f'sudo docker exec {DEFAULT_DOCKER_CONTAINER_NAME} whoami',
+        stream_logs=False,
+        require_outputs=True)
+    assert whoami_returncode == 0, (
+        f'Failed to get docker container user. Return '
+        f'code: {whoami_returncode}, Error: {whoami_stderr}')
+    docker_user = whoami_stdout.strip()
+    logger.debug(f'Docker container user: {docker_user}')
+    # Change host ssh port to 10022 to avoid conflict with docker.
+    # Docker is running with --net=host, which means the container
+    # will have the same IP address as the host machine. If both
+    # container and host sshd are running on the same port, the
+    # docker sshd will fail to start. So we change the host ssh
+    # port to 10022 to avoid conflict with docker, in the same time
+    # we restart the docker sshd service.
+    docker_host_ssh_setup_commands = [
+        'sudo sed -i "s/#Port 22/Port 10022/" /etc/ssh/sshd_config',
+        'sudo systemctl restart sshd',
+        # Restart ssh service in docker here since previous default
+        # ssh port is occupied by the host.
+        f'sudo docker exec {DEFAULT_DOCKER_CONTAINER_NAME} '
+        'bash --login -c -i "sudo service ssh restart"'
+    ]
+    runner.run('; '.join(docker_host_ssh_setup_commands), stream_logs=False)
+    return docker_user
+
+
 def build_dockerimage(task: task_mod.Task,
                       tag: str) -> Tuple[str, Dict[str, str]]:
     """
@@ -161,6 +227,8 @@ def build_dockerimage(task: task_mod.Task,
     if callable(task.run):
         raise ValueError(
             'Cannot build docker image for a task.run with function.')
+    assert task.docker_image is not None, 'Docker image cannot be None - ' \
+                                            'have you specified a task name? '
     _, img_metadata = create_dockerfile(base_image=task.docker_image,
                                         setup_command=task.setup,
                                         copy_path=f'{SKY_DOCKER_WORKDIR}/',
@@ -231,3 +299,142 @@ def bash_codegen(workdir_name: str,
         with open(out_path, 'w') as fp:
             fp.write(script_contents)
     return script_contents
+
+
+class SkyDockerCommandRunner(DockerCommandRunner):
+    """A DockerCommandRunner that install rsync before running init.
+
+    The code is borowed from
+    `ray.autoscaler._private.command_runner.DockerCommandRunner`."""
+
+    def run_init(self, *, as_head: bool, file_mounts: Dict[str, str],
+                 sync_run_yet: bool):
+        bootstrap_mounts = [
+            '~/ray_bootstrap_config.yaml', '~/ray_bootstrap_key.pem'
+        ]
+
+        specific_image = self.docker_config.get(
+            f'{"head" if as_head else "worker"}_image',
+            self.docker_config.get('image'))
+
+        self._check_docker_installed()
+        if self.docker_config.get('pull_before_run', True):
+            assert specific_image, ('Image must be included in config if ' +
+                                    'pull_before_run is specified')
+            self.run('{} pull {}'.format(self.docker_cmd, specific_image),
+                     run_env='host')
+        else:
+
+            self.run(f'{self.docker_cmd} image inspect {specific_image} '
+                     '1> /dev/null  2>&1 || '
+                     f'{self.docker_cmd} pull {specific_image}')
+
+        # Bootstrap files cannot be bind mounted because docker opens the
+        # underlying inode. When the file is switched, docker becomes outdated.
+        cleaned_bind_mounts = file_mounts.copy()
+        for mnt in bootstrap_mounts:
+            cleaned_bind_mounts.pop(mnt, None)
+
+        docker_run_executed = False
+
+        container_running = self._check_container_status()
+        requires_re_init = False
+        if container_running:
+            requires_re_init = self._check_if_container_restart_is_needed(
+                specific_image, cleaned_bind_mounts)
+            if requires_re_init:
+                self.run(f'{self.docker_cmd} stop {self.container_name}',
+                         run_env='host')
+
+        if (not container_running) or requires_re_init:
+            if not sync_run_yet:
+                # Do not start the actual image as we need to run file_sync
+                # first to ensure that all folders are created with the
+                # correct ownership. Docker will create the folders with
+                # `root` as the owner.
+                return True
+            # Get home directory
+            image_env = (self.ssh_command_runner.run(
+                f'{self.docker_cmd} ' + 'inspect -f "{{json .Config.Env}}" ' +
+                specific_image,
+                with_output=True,
+            ).decode().strip())
+            home_directory = '/root'
+            try:
+                for env_var in json.loads(image_env):
+                    if env_var.startswith('HOME='):
+                        home_directory = env_var.split('HOME=')[1]
+                        break
+            except json.JSONDecodeError as e:
+                cli_logger.error(
+                    'Unable to deserialize `image_env` to Python object. '
+                    f'The `image_env` is:\n{image_env}')
+                raise e
+
+            user_docker_run_options = self.docker_config.get(
+                'run_options', []) + self.docker_config.get(
+                    f'{"head" if as_head else "worker"}_run_options', [])
+            start_command = docker_start_cmds(
+                self.ssh_command_runner.ssh_user,
+                specific_image,
+                cleaned_bind_mounts,
+                self.container_name,
+                self._configure_runtime(
+                    self._auto_configure_shm(user_docker_run_options)),
+                self.ssh_command_runner.cluster_name,
+                home_directory,
+                self.docker_cmd,
+            )
+            self.run(start_command, run_env='host')
+            docker_run_executed = True
+
+        self.run('apt-get update; apt-get install -y rsync')
+
+        # Explicitly copy in ray bootstrap files.
+        for mount in bootstrap_mounts:
+            if mount in file_mounts:
+                if not sync_run_yet:
+                    # NOTE(ilr) This rsync is needed because when starting from
+                    #  a stopped instance,  /tmp may be deleted and `run_init`
+                    # is called before the first `file_sync` happens
+                    self.run_rsync_up(file_mounts[mount], mount)
+                self.ssh_command_runner.run(
+                    'rsync -e "{cmd} exec -i" -avz {src} {container}:{dst}'.
+                    format(
+                        cmd=self.docker_cmd,
+                        src=os.path.join(
+                            self._get_docker_host_mount_location(
+                                self.ssh_command_runner.cluster_name),
+                            mount,
+                        ),
+                        container=self.container_name,
+                        dst=self._docker_expand_user(mount),
+                    ))
+                try:
+                    # Check if the current user has read permission.
+                    # If they do not, try to change ownership!
+                    self.run(f'cat {mount} >/dev/null 2>&1 || '
+                             f'sudo chown $(id -u):$(id -g) {mount}')
+                # pylint: disable=broad-except
+                except Exception:
+                    lsl_string = (self.run(
+                        f'ls -l {mount}',
+                        with_output=True).decode('utf-8').strip())
+                    # The string is of format <Permission> <Links>
+                    # <Owner> <Group> <Size> <Date> <Name>
+                    permissions = lsl_string.split(' ')[0]
+                    owner = lsl_string.split(' ')[2]
+                    group = lsl_string.split(' ')[3]
+                    current_user = (self.run(
+                        'whoami', with_output=True).decode('utf-8').strip())
+                    cli_logger.warning(
+                        f'File ({mount}) is owned by user:{owner} and group:'
+                        f'{group} with permissions ({permissions}). The '
+                        f'current user ({current_user}) does not have '
+                        'permission to read these files, and Ray may not be '
+                        'able to autoscale. This can be resolved by '
+                        'installing `sudo` in your container, or adding a '
+                        f'command like "chown {current_user} {mount}" to '
+                        'your `setup_commands`.')
+        self.initialized = True
+        return docker_run_executed
diff --git a/sky/backends/onprem_utils.py b/sky/backends/onprem_utils.py
index 0666210e5aa..b666e2770d9 100644
--- a/sky/backends/onprem_utils.py
+++ b/sky/backends/onprem_utils.py
@@ -140,8 +140,10 @@ def get_python_executable(cluster_name: str) -> str:
     return config['python']
 
 
-def get_job_owner(cluster_yaml: str) -> str:
+def get_job_owner(cluster_yaml: str, docker_user: Optional[str] = None) -> str:
     """Get the owner of the job."""
+    if docker_user:
+        return docker_user
     cluster_config = common_utils.read_yaml(os.path.expanduser(cluster_yaml))
     # User name is guaranteed to exist (on all jinja files)
     return cluster_config['auth']['ssh_user']
diff --git a/sky/benchmark/benchmark_utils.py b/sky/benchmark/benchmark_utils.py
index 2522ea88ca1..a189de0f8c2 100644
--- a/sky/benchmark/benchmark_utils.py
+++ b/sky/benchmark/benchmark_utils.py
@@ -94,9 +94,14 @@ def _get_optimized_resources(
     for config in candidate_configs:
         with sky.Dag() as dag:
             resources = config.get('resources', None)
+            # pylint: disable=import-outside-toplevel
+            from sky.backends.docker_utils import \
+                extract_docker_image_from_resources
+            docker_image = extract_docker_image_from_resources(resources)
             resources = sky.Resources.from_yaml_config(resources)
             task = sky.Task()
             task.set_resources({resources})
+            task.docker_image = docker_image
 
         dag = sky.optimize(dag, quiet=True)
         task = dag.tasks[0]
diff --git a/sky/cli.py b/sky/cli.py
index 9f9768fcef8..5a92b02ec46 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -1016,6 +1016,10 @@ def _make_task_from_entrypoint_with_overrides(
     if workdir is not None:
         task.workdir = workdir
 
+    if image_id and image_id.startswith('docker:'):
+        task.docker_image = image_id[len('docker:'):]
+        image_id = None
+
     if onprem_utils.check_local_cloud_args(cloud, cluster, yaml_config):
         cloud = 'local'
 
diff --git a/sky/clouds/azure.py b/sky/clouds/azure.py
index b5fb32ee924..27f31af77ec 100644
--- a/sky/clouds/azure.py
+++ b/sky/clouds/azure.py
@@ -1,9 +1,11 @@
 """Azure."""
+import base64
 import functools
 import json
 import os
 import re
 import subprocess
+import textwrap
 import typing
 from typing import Dict, Iterator, List, Optional, Tuple
 
@@ -233,6 +235,18 @@ def make_deploy_resources_variables(
         gen_version = azure_catalog.get_gen_version_from_instance_type(
             r.instance_type)
         image_config = self._get_image_config(gen_version, r.instance_type)
+        # Setup commands to eliminate the banner and restart sshd.
+        # This script will modify /etc/ssh/sshd_config and add a bash script
+        # into .bashrc. The bash script will restart sshd if it has not been
+        # restarted, identified by a file /tmp/__restarted is existing.
+        # pylint: disable=line-too-long
+        cloud_init_setup_commands = base64.b64encode(
+            textwrap.dedent("""\
+            #cloud-config
+                runcmd:
+                - sed -i 's/#Banner none/Banner none/' /etc/ssh/sshd_config
+                - echo '\\nif [ ! -f "/tmp/__restarted" ]; then\\n  sudo systemctl restart ssh\\n  sleep 2\\n  touch /tmp/__restarted\\nfi' >> /home/azureuser/.bashrc
+            """).encode('utf-8')).decode('utf-8')
         return {
             'instance_type': r.instance_type,
             'custom_resources': custom_resources,
@@ -241,7 +255,8 @@ def make_deploy_resources_variables(
             # Azure does not support specific zones.
             'zones': None,
             **image_config,
-            'disk_tier': Azure._get_disk_type(r.disk_tier)
+            'disk_tier': Azure._get_disk_type(r.disk_tier),
+            'cloud_init_setup_commands': cloud_init_setup_commands
         }
 
     def get_feasible_launchable_resources(self, resources):
diff --git a/sky/skylet/log_lib.py b/sky/skylet/log_lib.py
index 45e82cc4671..4ad3d2efddb 100644
--- a/sky/skylet/log_lib.py
+++ b/sky/skylet/log_lib.py
@@ -242,6 +242,8 @@ def run_with_log(
 
 def make_task_bash_script(codegen: str,
                           env_vars: Optional[Dict[str, str]] = None) -> str:
+    # set -m to enable job control.
+    # Reference: https://stackoverflow.com/questions/11821378/what-does-bashno-job-control-in-this-shell-mean # pylint: disable=line-too-long
     # set -a is used for exporting all variables functions to the environment
     # so that bash `user_script` can access `conda activate`. Detail: #436.
     # Reference: https://www.gnu.org/software/bash/manual/html_node/The-Set-Builtin.html # pylint: disable=line-too-long
@@ -249,6 +251,7 @@ def make_task_bash_script(codegen: str,
     script = [
         textwrap.dedent(f"""\
             #!/bin/bash
+            set -m
             source ~/.bashrc
             set -a
             . $(conda info --base 2> /dev/null)/etc/profile.d/conda.sh > /dev/null 2>&1 || true
diff --git a/sky/skylet/providers/aws/config.py b/sky/skylet/providers/aws/config.py
index 79ab8ebdc23..497d47184ee 100644
--- a/sky/skylet/providers/aws/config.py
+++ b/sky/skylet/providers/aws/config.py
@@ -1035,7 +1035,13 @@ def _create_default_ssh_inbound_rules():
             "ToPort": 22,
             "IpProtocol": "tcp",
             "IpRanges": [{"CidrIp": "0.0.0.0/0"}],
-        }
+        },
+        {
+            "FromPort": 10022,
+            "ToPort": 10022,
+            "IpProtocol": "tcp",
+            "IpRanges": [{"CidrIp": "0.0.0.0/0"}],
+        },
     ]
 
 
diff --git a/sky/skylet/providers/aws/node_provider.py b/sky/skylet/providers/aws/node_provider.py
index 86297b748f7..44cf5f2c7ed 100644
--- a/sky/skylet/providers/aws/node_provider.py
+++ b/sky/skylet/providers/aws/node_provider.py
@@ -24,7 +24,9 @@
     resource_cache,
     client_cache,
 )
+from sky.backends.docker_utils import SkyDockerCommandRunner
 from ray.autoscaler._private.cli_logger import cli_logger, cf
+from ray.autoscaler._private.command_runner import SSHCommandRunner
 from ray.autoscaler._private.constants import BOTO_MAX_RETRIES, BOTO_CREATE_MAX_RETRIES
 from ray.autoscaler._private.log_timer import LogTimer
 from ray.autoscaler.node_provider import NodeProvider
@@ -714,6 +716,30 @@ def fillout_available_node_types_resources(
                 )
         return cluster_config
 
+    def get_command_runner(
+        self,
+        log_prefix,
+        node_id,
+        auth_config,
+        cluster_name,
+        process_runner,
+        use_internal_ip,
+        docker_config=None,
+    ):
+        common_args = {
+            "log_prefix": log_prefix,
+            "node_id": node_id,
+            "provider": self,
+            "auth_config": auth_config,
+            "cluster_name": cluster_name,
+            "process_runner": process_runner,
+            "use_internal_ip": use_internal_ip,
+        }
+        if docker_config and docker_config["container_name"] != "":
+            return SkyDockerCommandRunner(docker_config, **common_args)
+        else:
+            return SSHCommandRunner(**common_args)
+
 
 class AWSNodeProviderV2(AWSNodeProvider):
     """Same as V1, except head and workers use a SkyPilot IAM role.
diff --git a/sky/skylet/providers/azure/azure-config-template.json b/sky/skylet/providers/azure/azure-config-template.json
index 1a13a67a121..8b26c9be9d6 100644
--- a/sky/skylet/providers/azure/azure-config-template.json
+++ b/sky/skylet/providers/azure/azure-config-template.json
@@ -26,7 +26,7 @@
         "subnetName": "[concat('ray-', parameters('clusterId'), '-subnet')]"
     },
     "resources": [
-       {
+        {
             "type": "Microsoft.ManagedIdentity/userAssignedIdentities",
             "apiVersion": "2018-11-30",
             "location": "[variables('location')]",
@@ -56,7 +56,7 @@
                     {
                         "name": "SSH",
                         "properties": {
-                            "priority": 1000,
+                            "priority": 1001,
                             "protocol": "TCP",
                             "access": "Allow",
                             "direction": "Inbound",
@@ -65,6 +65,19 @@
                             "destinationAddressPrefix": "*",
                             "destinationPortRange": "22"
                         }
+                    },
+                    {
+                        "name": "SSH-DOCKER-HOST",
+                        "properties": {
+                            "priority": 1000,
+                            "protocol": "TCP",
+                            "access": "Allow",
+                            "direction": "Inbound",
+                            "sourceAddressPrefix": "*",
+                            "sourcePortRange": "*",
+                            "destinationAddressPrefix": "*",
+                            "destinationPortRange": "10022"
+                        }
                     }
                 ]
             }
diff --git a/sky/skylet/providers/azure/azure-vm-template.json b/sky/skylet/providers/azure/azure-vm-template.json
index 03ca08944c2..52e82dc532c 100644
--- a/sky/skylet/providers/azure/azure-vm-template.json
+++ b/sky/skylet/providers/azure/azure-vm-template.json
@@ -117,6 +117,12 @@
             "metadata": {
                 "description": "OS disk tier."
             }
+        },
+        "cloudInitSetupCommands": {
+            "type": "string",
+            "metadata": {
+                "description": "Base64 encoded cloud-init setup commands."
+            }
         }
     },
     "variables": {
@@ -260,7 +266,8 @@
                                 }
                             ]
                         }
-                    }
+                    },
+                    "customData": "[parameters('cloudInitSetupCommands')]"
                 },
                 "priority": "[parameters('priority')]",
                 "billingProfile": "[parameters('billingProfile')]"
diff --git a/sky/skylet/providers/azure/node_provider.py b/sky/skylet/providers/azure/node_provider.py
index 1a0306e5b67..4908704db57 100644
--- a/sky/skylet/providers/azure/node_provider.py
+++ b/sky/skylet/providers/azure/node_provider.py
@@ -15,6 +15,8 @@
     bootstrap_azure,
     get_azure_sdk_function,
 )
+from sky.backends.docker_utils import SkyDockerCommandRunner
+from ray.autoscaler._private.command_runner import SSHCommandRunner
 from ray.autoscaler.node_provider import NodeProvider
 from ray.autoscaler.tags import (
     TAG_RAY_CLUSTER_NAME,
@@ -419,3 +421,27 @@ def _get_cached_node(self, node_id):
     @staticmethod
     def bootstrap_config(cluster_config):
         return bootstrap_azure(cluster_config)
+
+    def get_command_runner(
+        self,
+        log_prefix,
+        node_id,
+        auth_config,
+        cluster_name,
+        process_runner,
+        use_internal_ip,
+        docker_config=None,
+    ):
+        common_args = {
+            "log_prefix": log_prefix,
+            "node_id": node_id,
+            "provider": self,
+            "auth_config": auth_config,
+            "cluster_name": cluster_name,
+            "process_runner": process_runner,
+            "use_internal_ip": use_internal_ip,
+        }
+        if docker_config and docker_config["container_name"] != "":
+            return SkyDockerCommandRunner(docker_config, **common_args)
+        else:
+            return SSHCommandRunner(**common_args)
diff --git a/sky/skylet/providers/gcp/config.py b/sky/skylet/providers/gcp/config.py
index 7d5d9c94251..a9369f233e0 100644
--- a/sky/skylet/providers/gcp/config.py
+++ b/sky/skylet/providers/gcp/config.py
@@ -550,6 +550,10 @@ def _merge_and_refine_rule(rules):
         source2rules: Dict[Tuple[str, str], Dict[str, Set[int]]] = {}
         source2allowed_list: Dict[Tuple[str, str], List[Dict[str, str]]] = {}
         for rule in rules:
+            # Disable the rule if it is not apply to all VMs
+            # See https://developers.google.com/resources/api-libraries/documentation/compute/alpha/python/latest/compute_alpha.networks.html#getEffectiveFirewalls # pylint: disable=line-too-long
+            if rule.get("targetTags", None) is not None:
+                continue
             direction = rule.get("direction", "")
             sources = rule.get("sourceRanges", [])
             allowed = rule.get("allowed", [])
diff --git a/sky/skylet/providers/gcp/constants.py b/sky/skylet/providers/gcp/constants.py
index 282d9b40885..9ccc3865a87 100644
--- a/sky/skylet/providers/gcp/constants.py
+++ b/sky/skylet/providers/gcp/constants.py
@@ -31,6 +31,17 @@
         ],
         "sourceRanges": ["0.0.0.0/0"],
     },
+    # Allow ssh connection to docker host from anywhere.
+    {
+        "direction": "INGRESS",
+        "allowed": [
+            {
+                "IPProtocol": "tcp",
+                "ports": ["10022"],
+            }
+        ],
+        "sourceRanges": ["0.0.0.0/0"],
+    },
 ]
 # Template when creating firewall rules for a new VPC.
 FIREWALL_RULES_TEMPLATE = [
@@ -63,6 +74,21 @@
         ],
         "sourceRanges": ["0.0.0.0/0"],
     },
+    {
+        "name": "{VPC_NAME}-allow-docker-ssh",
+        "description": "Allows TCP connections from any source to docker host inside any instance on the network using port 10022.",
+        "network": "projects/{PROJ_ID}/global/networks/{VPC_NAME}",
+        "selfLink": "projects/{PROJ_ID}/global/firewalls/{VPC_NAME}-allow-docker-ssh",
+        "direction": "INGRESS",
+        "priority": 65534,
+        "allowed": [
+            {
+                "IPProtocol": "tcp",
+                "ports": ["10022"],
+            }
+        ],
+        "sourceRanges": ["0.0.0.0/0"],
+    },
     {
         "name": "{VPC_NAME}-allow-icmp",
         "description": "Allows ICMP connections from any source to any instance on the network.",
diff --git a/sky/skylet/providers/gcp/node_provider.py b/sky/skylet/providers/gcp/node_provider.py
index 1a3443f52e6..403e4d41615 100644
--- a/sky/skylet/providers/gcp/node_provider.py
+++ b/sky/skylet/providers/gcp/node_provider.py
@@ -12,6 +12,7 @@
     construct_clients_from_provider_config,
     get_node_type,
 )
+from sky.backends.docker_utils import SkyDockerCommandRunner
 
 from ray.autoscaler.tags import (
     TAG_RAY_LAUNCH_CONFIG,
@@ -19,6 +20,7 @@
     TAG_RAY_USER_NODE_TYPE,
 )
 from ray.autoscaler._private.cli_logger import cf, cli_logger
+from ray.autoscaler._private.command_runner import SSHCommandRunner
 
 
 # The logic has been abstracted away here to allow for different GCP resources
@@ -362,3 +364,27 @@ def _get_cached_node(self, node_id: str) -> GCPNode:
     @staticmethod
     def bootstrap_config(cluster_config):
         return bootstrap_gcp(cluster_config)
+
+    def get_command_runner(
+        self,
+        log_prefix,
+        node_id,
+        auth_config,
+        cluster_name,
+        process_runner,
+        use_internal_ip,
+        docker_config=None,
+    ):
+        common_args = {
+            "log_prefix": log_prefix,
+            "node_id": node_id,
+            "provider": self,
+            "auth_config": auth_config,
+            "cluster_name": cluster_name,
+            "process_runner": process_runner,
+            "use_internal_ip": use_internal_ip,
+        }
+        if docker_config and docker_config["container_name"] != "":
+            return SkyDockerCommandRunner(docker_config, **common_args)
+        else:
+            return SSHCommandRunner(**common_args)
diff --git a/sky/skylet/providers/ibm/vpc_provider.py b/sky/skylet/providers/ibm/vpc_provider.py
index ef6c9147bac..b3d6c9fac05 100644
--- a/sky/skylet/providers/ibm/vpc_provider.py
+++ b/sky/skylet/providers/ibm/vpc_provider.py
@@ -21,6 +21,7 @@
     "outbound_udp_all": "selected security group is missing rule permitting outbound UDP access\n",
     "inbound_tcp_sg": "selected security group is missing rule permitting inbound tcp traffic inside selected security group\n",
     "inbound_tcp_22": "selected security group is missing rule permitting inbound traffic to tcp port 22 required for ssh\n",
+    "inbound_tcp_10022": "selected security group is missing rule permitting inbound traffic to tcp port 10022 required for ssh\n",
 }
 INSECURE_RULES = {
     "inbound_tcp_6379": "selected security group is missing rule permitting inbound traffic to tcp port 6379 required for Redis\n",
@@ -412,6 +413,7 @@ def get_unsatisfied_security_group_rules(self, sg_id):
                         # there a rule permitting all traffic
                         unsatisfied_rules.pop("inbound_tcp_sg", None)
                         unsatisfied_rules.pop("inbound_tcp_22", None)
+                        unsatisfied_rules.pop("inbound_tcp_10022", None)
                         unsatisfied_rules.pop("inbound_tcp_6379", None)
                         unsatisfied_rules.pop("inbound_tcp_8265", None)
 
@@ -420,6 +422,7 @@ def get_unsatisfied_security_group_rules(self, sg_id):
                             # all ports are open
                             unsatisfied_rules.pop("inbound_tcp_sg", None)
                             unsatisfied_rules.pop("inbound_tcp_22", None)
+                            unsatisfied_rules.pop("inbound_tcp_10022", None)
                             unsatisfied_rules.pop("inbound_tcp_6379", None)
                             unsatisfied_rules.pop("inbound_tcp_8265", None)
                         else:
@@ -427,6 +430,8 @@ def get_unsatisfied_security_group_rules(self, sg_id):
                             port_max = rule["port_max"]
                             if port_min <= 22 and port_max >= 22:
                                 unsatisfied_rules.pop("inbound_tcp_22", None)
+                            elif port_min <= 10022 and port_max >= 10022:
+                                unsatisfied_rules.pop("inbound_tcp_10022", None)
                             elif port_min <= 6379 and port_max >= 6379:
                                 unsatisfied_rules.pop("inbound_tcp_6379", None)
                             elif port_min <= 8265 and port_max >= 8265:
diff --git a/sky/task.py b/sky/task.py
index a8a41695c1e..86756056be1 100644
--- a/sky/task.py
+++ b/sky/task.py
@@ -133,9 +133,8 @@ def __init__(
             setup/run command, where ``run`` can either be a str, meaning all
             nodes get the same command, or a lambda, with the semantics
             documented above.
-          docker_image: (EXPERIMENTAL: Only in effect when LocalDockerBackend
-            is used.) The base docker image that this Task will be built on.
-            Defaults to 'gpuci/miniforge-cuda:11.4-devel-ubuntu18.04'.
+          docker_image: The base docker image that this Task will be built on.
+            If not specified, task won't run inside docker.
         """
         self.name = name
         self.run = run
@@ -145,8 +144,7 @@ def __init__(
         self.setup = setup
         self._envs = envs or {}
         self.workdir = workdir
-        self.docker_image = (docker_image if docker_image else
-                             'gpuci/miniforge-cuda:11.4-devel-ubuntu18.04')
+        self.docker_image = docker_image
         # Ignore type error due to a mypy bug.
         # https://github.com/python/mypy/issues/3004
         self.num_nodes = num_nodes  # type: ignore
@@ -327,9 +325,16 @@ def from_yaml(yaml_path: str) -> 'Task':
                              estimated_size_gigabytes=estimated_size_gigabytes)
 
         resources = config.pop('resources', None)
+
+        # pylint: disable=import-outside-toplevel
+        from sky.backends.docker_utils import \
+            extract_docker_image_from_resources
+        task.docker_image = extract_docker_image_from_resources(resources)
+
         resources = sky.Resources.from_yaml_config(resources)
 
         task.set_resources({resources})
+
         assert not config, f'Invalid task args: {config.keys()}'
         return task
 
@@ -872,4 +877,7 @@ def __repr__(self):
             s += f'\n  resources: {list(self.resources)[0]}'
         else:
             s += '\n  resources: default instances'
+
+        if self.docker_image is not None:
+            s += f'\n  docker_image: {self.docker_image}'
         return s
diff --git a/sky/templates/aws-ray.yml.j2 b/sky/templates/aws-ray.yml.j2
index 11115723a1c..4da7f118a90 100644
--- a/sky/templates/aws-ray.yml.j2
+++ b/sky/templates/aws-ray.yml.j2
@@ -5,6 +5,18 @@ max_workers: {{num_nodes - 1}}
 upscaling_speed: {{num_nodes - 1}}
 idle_timeout_minutes: 60
 
+{%- if use_docker %}
+docker:
+  image: {{docker_image}}
+  container_name: {{docker_container_name}}
+  run_options:
+    {%- if custom_resources is none %}
+    []
+    {%- else %}
+    - --gpus all
+    {%- endif %}
+{%- endif %}
+
 provider:
   type: external
   module: sky.skylet.providers.aws.AWSNodeProviderV2
@@ -172,6 +184,7 @@ initialization_commands: []
 # Increment the following for catching performance bugs easier:
 #   current num items (num SSH connections): 1
 setup_commands:
+  # Line echo 'export DEBIAN_FRONTEND=noninteractive': Disable apt-get from asking user input during installation. https://askubuntu.com/questions/909277/avoiding-user-interaction-with-tzdata-when-installing-certbot-in-a-docker-contai
   # Disable `unattended-upgrades` to prevent apt-get from hanging. It should be called at the beginning before the process started to avoid being blocked. (This is a temporary fix.)
   # Create ~/.ssh/config file in case the file does not exist in the custom image.
   # Make sure python3 & pip3 are available on this image.
@@ -181,19 +194,34 @@ setup_commands:
   # Line 'sudo grep ..': set the number of threads per process to unlimited to avoid ray job submit stucking issue when the number of running ray jobs increase.
   # Line 'mkdir -p ..': disable host key check
   # Line 'python3 -c ..': patch the buggy ray files and enable `-o allow_other` option for `goofys`
+  {%- if use_docker %}
+  - echo '[ "$(whoami)" == "root" ] && alias sudo=""' >> ~/.bashrc;
+    echo 'export DEBIAN_FRONTEND=noninteractive' >> ~/.bashrc;
+  {%- endif %}
   - mkdir -p ~/.ssh; touch ~/.ssh/config;
+    {%- if use_docker %}
+    source ~/.bashrc;
+    sudo apt-get update;
+    sudo apt-get install -y curl wget patch openssh-server;
+    {%- endif %}
+    (which conda > /dev/null 2>&1 && conda init > /dev/null) || (wget -nc https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && bash Miniconda3-latest-Linux-x86_64.sh -b && eval "$(~/miniconda3/bin/conda shell.bash hook)" && conda init && conda config --set auto_activate_base true);
+    source ~/.bashrc;
     pip3 --version > /dev/null 2>&1 || (curl -sSL https://bootstrap.pypa.io/get-pip.py -o get-pip.py && python3 get-pip.py && echo "PATH=$HOME/.local/bin:$PATH" >> ~/.bashrc);
     (type -a python | grep -q python3) || echo 'alias python=python3' >> ~/.bashrc;
     (type -a pip | grep -q pip3) || echo 'alias pip=pip3' >> ~/.bashrc;
-    (which conda > /dev/null 2>&1 && conda init > /dev/null) || (wget -nc https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && bash Miniconda3-latest-Linux-x86_64.sh -b && eval "$(~/miniconda3/bin/conda shell.bash hook)" && conda init && conda config --set auto_activate_base true);
-    source ~/.bashrc;
     (pip3 list | grep "ray " | grep {{ray_version}} 2>&1 > /dev/null || pip3 install --exists-action w -U ray[default]=={{ray_version}}) && mkdir -p ~/sky_workdir && mkdir -p ~/.sky/sky_app;
     (pip3 list | grep "skypilot " && [ "$(cat {{sky_remote_path}}/current_sky_wheel_hash)" == "{{sky_wheel_hash}}" ]) || (pip3 uninstall skypilot -y; pip3 install "$(echo {{sky_remote_path}}/{{sky_wheel_hash}}/skypilot-{{sky_version}}*.whl)[aws]" && echo "{{sky_wheel_hash}}" > {{sky_remote_path}}/current_sky_wheel_hash || exit 1);
     sudo bash -c 'rm -rf /etc/security/limits.d; echo "* soft nofile 1048576" >> /etc/security/limits.conf; echo "* hard nofile 1048576" >> /etc/security/limits.conf';
+    {%- if not use_docker %}
     sudo grep -e '^DefaultTasksMax' /etc/systemd/system.conf || (sudo bash -c 'echo "DefaultTasksMax=infinity" >> /etc/systemd/system.conf'); sudo systemctl set-property user-$(id -u $(whoami)).slice TasksMax=infinity; sudo systemctl daemon-reload;
+    {%- endif %}
     mkdir -p ~/.ssh; (grep -Pzo -q "Host \*\n  StrictHostKeyChecking no" ~/.ssh/config) || printf "Host *\n  StrictHostKeyChecking no\n" >> ~/.ssh/config;
     python3 -c "from sky.skylet.ray_patches import patch; patch()" || exit 1;
-    [ -f /etc/fuse.conf ] && sudo sed -i 's/#user_allow_other/user_allow_other/g' /etc/fuse.conf || (sudo sh -c 'echo "user_allow_other" > /etc/fuse.conf'); # This is needed for `-o allow_other` option for `goofys`;
+    [ -f /etc/fuse.conf ] && sudo sed -i 's/#user_allow_other/user_allow_other/g' /etc/fuse.conf || (sudo sh -c 'echo "user_allow_other" > /etc/fuse.conf');
+    {%- if use_docker %}
+    echo "{{public_key}}" >> ~/.ssh/authorized_keys;
+    sudo sed -i 's/mesg n/tty -s \&\& mesg n/' ~/.profile;
+    {%- endif %}
 
 # Command to start ray on the head node. You don't need to change this.
 # NOTE: these are very performance-sensitive. Each new item opens/closes an SSH
diff --git a/sky/templates/azure-ray.yml.j2 b/sky/templates/azure-ray.yml.j2
index a489bf6d1ec..8bce4594fcd 100644
--- a/sky/templates/azure-ray.yml.j2
+++ b/sky/templates/azure-ray.yml.j2
@@ -5,6 +5,18 @@ max_workers: {{num_nodes - 1}}
 upscaling_speed: {{num_nodes - 1}}
 idle_timeout_minutes: 60
 
+{%- if use_docker %}
+docker:
+  image: {{docker_image}}
+  container_name: {{docker_container_name}}
+  run_options:
+    {%- if custom_resources is none %}
+    []
+    {%- else %}
+    - --gpus all
+    {%- endif %}
+{%- endif %}
+
 provider:
     type: external
     module: sky.skylet.providers.azure.AzureNodeProvider
@@ -46,6 +58,7 @@ available_node_types:
             imageVersion: {{image_version}}
             osDiskSizeGB: {{disk_size}}
             osDiskTier: {{disk_tier}}
+            cloudInitSetupCommands: {{cloud_init_setup_commands}}
             # optionally set priority to use Spot instances
             {%- if use_spot %}
             priority: Spot
@@ -71,6 +84,7 @@ available_node_types:
             imageVersion: {{image_version}}
             osDiskSizeGB: {{disk_size}}
             osDiskTier: {{disk_tier}}
+            cloudInitSetupCommands: {{cloud_init_setup_commands}}
           {%- if use_spot %}
             priority: Spot
             # set a maximum price for spot instances if desired
@@ -102,6 +116,7 @@ initialization_commands: []
 # Increment the following for catching performance bugs easier:
 #   current num items (num SSH connections): 1
 setup_commands:
+  # Line echo 'export DEBIAN_FRONTEND=noninteractive': Disable apt-get from asking user input during installation. https://askubuntu.com/questions/909277/avoiding-user-interaction-with-tzdata-when-installing-certbot-in-a-docker-contai
   # Disable `unattended-upgrades` to prevent apt-get from hanging. It should be called at the beginning before the process started to avoid being blocked. (This is a temporary fix.)
   # Create ~/.ssh/config file in case the file does not exist in the image.
   # Make sure python3 & pip3 are available on this image.
@@ -110,10 +125,20 @@ setup_commands:
   # Line 'mkdir -p ..': disable host key check
   # Line 'python3 -c ..': patch the buggy ray files and enable `-o allow_other` option for `goofys`
   # This also kills the service that is holding the lock on dpkg (problem only exists on aws/azure, not gcp)
+  {%- if use_docker %}
+  - echo '[ "$(whoami)" == "root" ] && alias sudo=""' >> ~/.bashrc;
+    echo 'export DEBIAN_FRONTEND=noninteractive' >> ~/.bashrc;
+  {%- endif %}
   - function mylsof { p=$(for pid in /proc/{0..9}*; do i=$(basename "$pid"); for file in "$pid"/fd/*; do link=$(readlink -e "$file"); if [ "$link" = "$1" ]; then echo "$i"; fi; done; done); echo "$p"; };
+    {%- if use_docker %}
+    source ~/.bashrc;
+    sudo apt-get update;
+    sudo apt-get install -y curl wget patch openssh-server;
+    {%- else %}
     sudo systemctl stop unattended-upgrades || true;
     sudo systemctl disable unattended-upgrades || true;
     sudo sed -i 's/Unattended-Upgrade "1"/Unattended-Upgrade "0"/g' /etc/apt/apt.conf.d/20auto-upgrades || true;
+    {%- endif %}
     p=$(mylsof "/var/lib/dpkg/lock-frontend"); echo "$p";
     sudo kill -9 `echo "$p" | tail -n 1` || true;
     sudo rm /var/lib/dpkg/lock-frontend;
@@ -121,18 +146,24 @@ setup_commands:
     sudo pkill -9 apt-get;
     sudo dpkg --configure --force-overwrite -a;
     mkdir -p ~/.ssh; touch ~/.ssh/config;
+    which conda > /dev/null 2>&1 || (wget -nc https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && bash Miniconda3-latest-Linux-x86_64.sh -b && eval "$(~/miniconda3/bin/conda shell.bash hook)" && conda init && conda config --set auto_activate_base true);
+    source ~/.bashrc;
     pip3 --version > /dev/null 2>&1 || (curl -sSL https://bootstrap.pypa.io/get-pip.py -o get-pip.py && python3 get-pip.py && echo "PATH=$HOME/.local/bin:$PATH" >> ~/.bashrc);
     (type -a python | grep -q python3) || echo 'alias python=python3' >> ~/.bashrc;
     (type -a pip | grep -q pip3) || echo 'alias pip=pip3' >> ~/.bashrc;
-    which conda > /dev/null 2>&1 || (wget -nc https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && bash Miniconda3-latest-Linux-x86_64.sh -b && eval "$(/home/azureuser/miniconda3/bin/conda shell.bash hook)" && conda init && conda config --set auto_activate_base true);
-    source ~/.bashrc;
     (pip3 list | grep "ray " | grep {{ray_version}} 2>&1 > /dev/null || pip3 install --exists-action w -U ray[default]=={{ray_version}}) && mkdir -p ~/sky_workdir && mkdir -p ~/.sky/sky_app && touch ~/.sudo_as_admin_successful;
     (pip3 list | grep "skypilot " && [ "$(cat {{sky_remote_path}}/current_sky_wheel_hash)" == "{{sky_wheel_hash}}" ]) || (pip3 uninstall skypilot -y; pip3 install "$(echo {{sky_remote_path}}/{{sky_wheel_hash}}/skypilot-{{sky_version}}*.whl)[azure]" && echo "{{sky_wheel_hash}}" > {{sky_remote_path}}/current_sky_wheel_hash || exit 1);
     sudo bash -c 'rm -rf /etc/security/limits.d; echo "* soft nofile 1048576" >> /etc/security/limits.conf; echo "* hard nofile 1048576" >> /etc/security/limits.conf';
+    {%- if not use_docker %}
     sudo grep -e '^DefaultTasksMax' /etc/systemd/system.conf || (sudo bash -c 'echo "DefaultTasksMax=infinity" >> /etc/systemd/system.conf'); sudo systemctl set-property user-$(id -u $(whoami)).slice TasksMax=infinity; sudo systemctl daemon-reload;
+    {%- endif %}
     mkdir -p ~/.ssh; (grep -Pzo -q "Host \*\n  StrictHostKeyChecking no" ~/.ssh/config) || printf "Host *\n  StrictHostKeyChecking no\n" >> ~/.ssh/config;
     python3 -c "from sky.skylet.ray_patches import patch; patch()" || exit 1;
     [ -f /etc/fuse.conf ] && sudo sed -i 's/#user_allow_other/user_allow_other/g' /etc/fuse.conf || (sudo sh -c 'echo "user_allow_other" > /etc/fuse.conf');
+    {%- if use_docker %}
+    echo "{{public_key}}" >> ~/.ssh/authorized_keys;
+    sudo sed -i 's/mesg n/tty -s \&\& mesg n/' ~/.profile;
+    {%- endif %}
 
 # Command to start ray on the head node. You don't need to change this.
 # NOTE: these are very performance-sensitive. Each new item opens/closes an SSH
diff --git a/sky/templates/gcp-ray.yml.j2 b/sky/templates/gcp-ray.yml.j2
index 3221ab2c21f..a205ccfaeda 100644
--- a/sky/templates/gcp-ray.yml.j2
+++ b/sky/templates/gcp-ray.yml.j2
@@ -6,6 +6,18 @@ max_workers: {{num_nodes - 1}}
 upscaling_speed: {{num_nodes - 1}}
 idle_timeout_minutes: 60
 
+{%- if use_docker %}
+docker:
+  image: {{docker_image}}
+  container_name: {{docker_container_name}}
+  run_options:
+    {%- if gpu is none %}
+    []
+    {%- else %}
+    - --gpus all
+    {%- endif %}
+{%- endif %}
+
 provider:
   # We use a custom node provider for GCP to support instance stop and reuse.
   type: external  # type: gcp
@@ -140,6 +152,7 @@ initialization_commands: []
 # Increment the following for catching performance bugs easier:
 #   current num items (num SSH connections): 1  (+1 if tpu_vm)
 setup_commands:
+  # Line echo 'export DEBIAN_FRONTEND=noninteractive': Disable apt-get from asking user input during installation. https://askubuntu.com/questions/909277/avoiding-user-interaction-with-tzdata-when-installing-certbot-in-a-docker-contai
   # Disable `unattended-upgrades` to prevent apt-get from hanging. It should be called at the beginning before the process started to avoid being blocked. (This is a temporary fix.)
   # Line 'mkdir -p ..': Create ~/.ssh/config file in case the file does not exist in the custom image.
   # Line 'pip3 --v ..': Make sure python3 & pip3 are available on this image.
@@ -150,10 +163,23 @@ setup_commands:
   # Line 'sudo grep ..': set the number of threads per process to unlimited to avoid ray job submit stucking issue when the number of running ray jobs increase.
   # Line 'mkdir -p ..': disable host key check
   # Line 'python3 -c ..': patch the buggy ray files and enable `-o allow_other` option for `goofys`
+  # Line 'curl ../libtpu.so': see https://cloud.google.com/tpu/docs/run-in-container
+  # Line 'sudo sed -i .. ~/.profile': to eliminate the error `mesg: ttyname failed: inappropriate ioctl for device`. https://www.educative.io/answers/error-mesg-ttyname-failed-inappropriate-ioctl-for-device
+  {%- if use_docker %}
+  - echo '[ "$(whoami)" == "root" ] && alias sudo=""' >> ~/.bashrc;
+    echo 'export DEBIAN_FRONTEND=noninteractive' >> ~/.bashrc;
+  {%- endif %}
   - function mylsof { p=$(for pid in /proc/{0..9}*; do i=$(basename "$pid"); for file in "$pid"/fd/*; do link=$(readlink -e "$file"); if [ "$link" = "$1" ]; then echo "$i"; fi; done; done); echo "$p"; };
+    p=$(mylsof "/var/lib/dpkg/lock-frontend"); echo "$p";
+    {%- if use_docker %}
+    source ~/.bashrc;
+    sudo apt-get update;
+    sudo apt-get install -y curl wget patch openssh-server;
+    {%- else %}
     sudo systemctl stop unattended-upgrades || true;
     sudo systemctl disable unattended-upgrades || true;
     sudo sed -i 's/Unattended-Upgrade "1"/Unattended-Upgrade "0"/g' /etc/apt/apt.conf.d/20auto-upgrades || true;
+    {%- endif %}
     p=$(mylsof "/var/lib/dpkg/lock-frontend"); echo "$p";
     sudo kill -9 `echo "$p" | tail -n 1` || true;
     sudo rm /var/lib/dpkg/lock-frontend;
@@ -161,22 +187,32 @@ setup_commands:
     sudo pkill -9 apt-get;
     sudo dpkg --configure --force-overwrite -a;
     mkdir -p ~/.ssh; touch ~/.ssh/config;
+    which conda > /dev/null 2>&1 || (wget -nc https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && bash Miniconda3-latest-Linux-x86_64.sh -b && eval "$(~/miniconda3/bin/conda shell.bash hook)" && conda init && conda config --set auto_activate_base true);
+    source ~/.bashrc;
     pip3 --version > /dev/null 2>&1 || (curl -sSL https://bootstrap.pypa.io/get-pip.py -o get-pip.py && python3 get-pip.py && echo "PATH=$HOME/.local/bin:$PATH" >> ~/.bashrc);
     (type -a python | grep -q python3) || echo 'alias python=python3' >> ~/.bashrc;
     (type -a pip | grep -q pip3) || echo 'alias pip=pip3' >> ~/.bashrc;
-    which conda > /dev/null 2>&1 || (wget -nc https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && bash Miniconda3-latest-Linux-x86_64.sh -b && eval "$(/home/gcpuser/miniconda3/bin/conda shell.bash hook)" && conda init && conda config --set auto_activate_base true);
-    source ~/.bashrc;
   {%- if tpu_vm %}
-    test -f /home/gcpuser/miniconda3/etc/profile.d/conda.sh && source /home/gcpuser/miniconda3/etc/profile.d/conda.sh && conda activate base || true;
+    test -f ~/miniconda3/etc/profile.d/conda.sh && source ~/miniconda3/etc/profile.d/conda.sh && conda activate base || true;
     pip3 install --upgrade google-api-python-client;
+    {%- if use_docker %}
+    curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.0/libtpu.so -o /lib/libtpu.so
+    {%- endif %}
   {%- endif %}
     (pip3 list | grep "ray " | grep {{ray_version}} 2>&1 > /dev/null || pip3 install --exists-action w -U ray[default]=={{ray_version}}) && mkdir -p ~/sky_workdir && mkdir -p ~/.sky/sky_app;
     (pip3 list | grep "skypilot " && [ "$(cat {{sky_remote_path}}/current_sky_wheel_hash)" == "{{sky_wheel_hash}}" ]) || (pip3 uninstall skypilot -y; pip3 install "$(echo {{sky_remote_path}}/{{sky_wheel_hash}}/skypilot-{{sky_version}}*.whl)[gcp]" && echo "{{sky_wheel_hash}}" > {{sky_remote_path}}/current_sky_wheel_hash || exit 1);
     sudo bash -c 'rm -rf /etc/security/limits.d; echo "* soft nofile 1048576" >> /etc/security/limits.conf; echo "* hard nofile 1048576" >> /etc/security/limits.conf';
+    {%- if not use_docker %}
     sudo grep -e '^DefaultTasksMax' /etc/systemd/system.conf || (sudo bash -c 'echo "DefaultTasksMax=infinity" >> /etc/systemd/system.conf'); sudo systemctl set-property user-$(id -u $(whoami)).slice TasksMax=infinity; sudo systemctl daemon-reload;
+    {%- endif %}
     mkdir -p ~/.ssh; (grep -Pzo -q "Host \*\n  StrictHostKeyChecking no" ~/.ssh/config) || printf "Host *\n  StrictHostKeyChecking no\n" >> ~/.ssh/config;
     python3 -c "from sky.skylet.ray_patches import patch; patch()" || exit 1;
     [ -f /etc/fuse.conf ] && sudo sed -i 's/#user_allow_other/user_allow_other/g' /etc/fuse.conf || (sudo sh -c 'echo "user_allow_other" > /etc/fuse.conf');
+    {%- if use_docker %}
+    echo "{{public_key}}" >> ~/.ssh/authorized_keys;
+    sudo sed -i 's/mesg n/tty -s \&\& mesg n/' ~/.profile;
+    {%- endif %}
+    
 
 # Command to start ray on the head node. You don't need to change this.
 # NOTE: these are very performance-sensitive. Each new item opens/closes an SSH


Code review comments:
diff_hunk: @@ -871,6 +878,10 @@ def write_cluster_config(
         f'open(os.path.expanduser("{constants.SKY_REMOTE_RAY_PORT_FILE}"), "w"))\''
     )
 
+    # pylint: disable=import-outside-toplevel
+    from sky.backends.docker_utils import \
+        DEFAULT_DOCKER_CONTAINER_NAME
path: sky/backends/backend_utils.py
text: @user1:
style: import module instead of variable:
```
from sky.backends import docker_utils
```

diff_hunk: @@ -777,6 +779,7 @@ def write_cluster_config(
         cluster_name: str,
         local_wheel_path: pathlib.Path,
         wheel_hash: str,
+        docker_image: Optional[str],
path: sky/backends/backend_utils.py
text: @user1:
Why do we need to pass the argument? We can get the image from the `cloud.make_deploy_resources_variables` function, right?

diff_hunk: @@ -1365,6 +1366,7 @@ def _retry_zones(
                     cluster_name,
                     self._local_wheel_path,
                     self._wheel_hash,
+                    docker_image,
path: sky/backends/cloud_vm_ray_backend.py
text: @user1:
Why do we need to pass the `docker_image` all the way down to here? Can we just use the `to_provision` to get the docker_image?

diff_hunk: @@ -225,23 +225,29 @@ def setup_gcp_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
             check=True,
             shell=True,
             stdout=subprocess.DEVNULL)
-        # Enable ssh port for all the instances
-        enable_ssh_cmd = ('gcloud compute firewall-rules create '
-                          'allow-ssh-ingress-from-iap '
-                          '--direction=INGRESS '
-                          '--action=allow '
-                          '--rules=tcp:22 '
-                          '--source-ranges=0.0.0.0/0')
-        proc = subprocess.run(enable_ssh_cmd,
-                              check=False,
-                              shell=True,
-                              stdout=subprocess.DEVNULL,
-                              stderr=subprocess.PIPE)
-        if proc.returncode != 0 and 'already exists' not in proc.stderr.decode(
-                'utf-8'):
-            subprocess_utils.handle_returncode(proc.returncode, enable_ssh_cmd,
-                                               'Failed to enable ssh port.',
-                                               proc.stderr.decode('utf-8'))
+
+        def allow_ssh_port(port):
+            # Enable ssh port for all the instances
+            enable_ssh_cmd = ('gcloud compute firewall-rules create '
+                              'allow-ssh-ingress-from-iap '
+                              '--direction=INGRESS '
+                              '--action=allow '
+                              f'--rules=tcp:{port} '
+                              '--source-ranges=0.0.0.0/0')
+            proc = subprocess.run(enable_ssh_cmd,
+                                  check=False,
+                                  shell=True,
+                                  stdout=subprocess.DEVNULL,
+                                  stderr=subprocess.PIPE)
+            if proc.returncode != 0 and \
+                'already exists' not in proc.stderr.decode('utf-8'):
+                subprocess_utils.handle_returncode(
+                    proc.returncode, enable_ssh_cmd,
+                    'Failed to enable ssh port.', proc.stderr.decode('utf-8'))
+
+        allow_ssh_port('22')
+        # For docker host machine
+        allow_ssh_port('10022')
path: sky/authentication.py
text: @user1:
Seems this only enables the ssh for 10022 when the oslogin is enabled. Is this intended? For projects without oslogin enabled, should we also enable the 10022 port?

diff_hunk: @@ -1016,6 +1016,10 @@ def _make_task_from_entrypoint_with_overrides(
     if workdir is not None:
         task.workdir = workdir
 
+    if image_id and image_id.startswith('docker:'):
+        task.docker_image = image_id[len('docker:'):]
+        image_id = None
path: sky/cli.py
text: @user1:
Why can't we just keep the `docker:` in the `image_id`, i.e. under the `task.resources`

diff_hunk: @@ -854,6 +857,10 @@ def write_cluster_config(
             assert region_name in ssh_proxy_command_config, (
                 region_name, ssh_proxy_command_config)
             ssh_proxy_command = ssh_proxy_command_config[region_name]
+    public_key_path = os.path.expanduser(auth.PUBLIC_SSH_KEY_PATH)
+    with open(public_key_path, 'r') as f:
+        public_key = f.read()
path: sky/backends/backend_utils.py
text: @user1:
It is a bit weird to see the public key retrieved in this file to set up the authentication. Is it possible to move this into the functions in the `authentication.py`?

diff_hunk: @@ -1106,9 +1124,18 @@ def wait_until_ray_cluster_ready(
             cluster_config_file, max_attempts=WAIT_HEAD_NODE_IP_MAX_ATTEMPTS)
     except RuntimeError as e:
         logger.error(e)
-        return False  # failed
+        return False, None  # failed
+
+    docker_user = None
+    if use_docker:
+        # pylint: disable=import-outside-toplevel
+        from sky.backends.docker_utils import docker_host_setup
+        docker_user = docker_host_setup(head_ip, cluster_config_file)
path: sky/backends/backend_utils.py
text: @user1:
style: import module instead of function
```suggestion
        from sky.backends import docker_utils
        docker_user = docker_utils.docker_host_setup(head_ip, cluster_config_file)
```

@user1:
Just curious, is it possible to set the user by the `run_options`, so we don't have to manage a different username than skypilot default one for VM?
<img width="544" alt="image" src="https://github.com/skypilot-org/skypilot/assets/6753189/f6e2ef28-dda7-4bcd-a872-cd5b8e207329">

The reason I mention this is that it looks not that clean to pass the `docker_user` around in the code.


@author:
Quick answer: this option requires the docker image containing a user named `my_username`, which is typically not the case for our default username `gcpuser` or `azureuser`. But this is the implementation previous to overriding Ray's `DockerCommandRunner`; since now we have the ability to modify the docker launch process, there may be some workarounds available and I'm trying it.

diff_hunk: @@ -327,9 +325,16 @@ def from_yaml(yaml_path: str) -> 'Task':
                              estimated_size_gigabytes=estimated_size_gigabytes)
 
         resources = config.pop('resources', None)
+
+        # pylint: disable=import-outside-toplevel
+        from sky.backends.docker_utils import \
+            extract_docker_image_from_resources
+        task.docker_image = extract_docker_image_from_resources(resources)
path: sky/task.py
text: @user1:
Why can't we just keep it as a part of the `resources` instead of a `task.docker_image`?