[Provisioner] Support docker in Lambda Cloud and TPU (#4115)

* [Provisioner] Support docker in Lambda Cloud

* fix permission issue

* merge with check docker installed

* add tpu support & test

* patch lambda cloud

* add comment

sky/clouds/azure.py

@@ -329,7 +329,6 @@ def make_deploy_resources_variables(
  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/skypilot:ssh_user/.bashrc
- - usermod -aG docker skypilot:ssh_user
  write_files:
  - path: /etc/apt/apt.conf.d/20auto-upgrades
  content: |

sky/clouds/gcp.py

@@ -477,6 +477,9 @@ def _failover_disk_tier() -> Optional[resources_utils.DiskTier]:
  'runtime_version']
  resources_vars['tpu_node_name'] = r.accelerator_args.get(
  'tpu_name')
+ # TPU VMs require privileged mode for docker containers to
+ # access TPU devices.
+ resources_vars['docker_run_options'] = ['--privileged']
  else:
  # Convert to GCP names:
  # https://cloud.google.com/compute/docs/gpus

sky/clouds/lambda_cloud.py

@@ -37,10 +37,6 @@ class Lambda(clouds.Cloud):
  _CLOUD_UNSUPPORTED_FEATURES = {
  clouds.CloudImplementationFeatures.STOP: 'Lambda cloud does not support stopping VMs.',
  clouds.CloudImplementationFeatures.CLONE_DISK_FROM_CLUSTER: f'Migrating disk is currently not supported on {_REPR}.',
- clouds.CloudImplementationFeatures.DOCKER_IMAGE: (
- f'Docker image is currently not supported on {_REPR}. '
- 'You can try running docker command inside the `run` section in task.yaml.'
- ),
  clouds.CloudImplementationFeatures.SPOT_INSTANCE: f'Spot instances are not supported in {_REPR}.',
  clouds.CloudImplementationFeatures.IMAGE_ID: f'Specifying image ID is not supported in {_REPR}.',
  clouds.CloudImplementationFeatures.CUSTOM_DISK_TIER: f'Custom disk tiers are not supported in {_REPR}.',
@@ -173,12 +169,20 @@ def make_deploy_resources_variables(
  else:
  custom_resources = None
 
- return {
+ resources_vars = {
  'instance_type': resources.instance_type,
  'custom_resources': custom_resources,
  'region': region.name,
  }
 
+ if acc_dict is not None:
+ # Lambda cloud's docker runtime information does not contain
+ # 'nvidia-container-runtime', causing no GPU option is added to
+ # the docker run command. We patch this by adding it here.
+ resources_vars['docker_run_options'] = ['--gpus all']
+
+ return resources_vars
+
  def _get_feasible_launchable_resources(
  self, resources: 'resources_lib.Resources'
  ) -> 'resources_utils.FeasibleResources':

sky/provision/docker_utils.py

@@ -253,12 +253,13 @@ def initialize(self) -> str:
  # issue with nvidia container toolkit:
  # https://github.com/NVIDIA/nvidia-container-toolkit/issues/48
  self._run(
- '[ -f /etc/docker/daemon.json ] || '
+ '{ which jq || sudo apt update && sudo apt install -y jq; } && '
+ '{ [ -f /etc/docker/daemon.json ] || '
  'echo "{}" | sudo tee /etc/docker/daemon.json;'
  'sudo jq \'.["exec-opts"] = ["native.cgroupdriver=cgroupfs"]\' '
  '/etc/docker/daemon.json > /tmp/daemon.json;'
  'sudo mv /tmp/daemon.json /etc/docker/daemon.json;'
- 'sudo systemctl restart docker')
+ 'sudo systemctl restart docker; } || true')
  user_docker_run_options = self.docker_config.get('run_options', [])
  start_command = docker_start_cmds(
  specific_image,
@@ -335,7 +336,11 @@ def initialize(self) -> str:
 
  def _check_docker_installed(self):
  no_exist = 'NoExist'
+ # SkyPilot: Add the current user to the docker group first (if needed),
+ # before checking if docker is installed to avoid permission issues.
  cleaned_output = self._run(
+ 'id -nG $USER | grep -qw docker || '
+ 'sudo usermod -aG docker $USER > /dev/null 2>&1;'
  f'command -v {self.docker_cmd} || echo {no_exist!r}')
  if no_exist in cleaned_output or 'docker' not in cleaned_output:
  logger.error(
@@ -424,8 +429,8 @@ def _auto_configure_shm(self, run_options: List[str]) -> List[str]:
  def _check_container_exited(self) -> bool:
  if self.initialized:
  return True
- output = (self._run(check_docker_running_cmd(self.container_name,
-  self.docker_cmd),
-  wait_for_docker_daemon=True))
- return 'false' in output.lower(
- ) and 'no such object' not in output.lower()
+ output = self._run(check_docker_running_cmd(self.container_name,
+ self.docker_cmd),
+ wait_for_docker_daemon=True)
+ return ('false' in output.lower() and
+  'no such object' not in output.lower())

sky/provision/paperspace/utils.py

@@ -132,6 +132,8 @@ def set_sky_key_script(self, public_key: str) -> None:
  'apt-get update \n'
  'apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin \n' # pylint: disable=line-too-long
  'fi \n'
+ # TODO(tian): Maybe remove this as well since we are now adding
+ # users to docker group in the DockerInitializer. Need to test.
  'usermod -aG docker paperspace \n'
  f'echo "{public_key}" >> /home/paperspace/.ssh/authorized_keys \n')
  try:

sky/resources.py

@@ -842,12 +842,6 @@ def _try_validate_image_id(self) -> None:
 
  if self.extract_docker_image() is not None:
  # TODO(tian): validate the docker image exists / of reasonable size
- if self.accelerators is not None:
- for acc in self.accelerators.keys():
- if acc.lower().startswith('tpu'):
- with ux_utils.print_exception_no_traceback():
- raise ValueError(
- 'Docker image is not supported for TPU VM.')
  if self.cloud is not None:
  self.cloud.check_features_are_supported(
  self, {clouds.CloudImplementationFeatures.DOCKER_IMAGE})
@@ -1032,25 +1026,30 @@ def make_deploy_variables(self, cluster_name: resources_utils.ClusterName,
  self.accelerators is not None):
  initial_setup_commands = [constants.DISABLE_GPU_ECC_COMMAND]
 
+ docker_image = self.extract_docker_image()
+
+ # Cloud specific variables
+ cloud_specific_variables = self.cloud.make_deploy_resources_variables(
+ self, cluster_name, region, zones, dryrun)
+
  # Docker run options
  docker_run_options = skypilot_config.get_nested(
  ('docker', 'run_options'),
  default_value=[],
  override_configs=self.cluster_config_overrides)
  if isinstance(docker_run_options, str):
  docker_run_options = [docker_run_options]
+ # Special accelerator runtime might require additional docker run
+ # options. e.g., for TPU, we need --privileged.
+ if 'docker_run_options' in cloud_specific_variables:
+ docker_run_options.extend(
+ cloud_specific_variables['docker_run_options'])
  if docker_run_options and isinstance(self.cloud, clouds.Kubernetes):
  logger.warning(
  f'{colorama.Style.DIM}Docker run options are specified, '
  'but ignored for Kubernetes: '
  f'{" ".join(docker_run_options)}'
  f'{colorama.Style.RESET_ALL}')
-
- docker_image = self.extract_docker_image()
-
- # Cloud specific variables
- cloud_specific_variables = self.cloud.make_deploy_resources_variables(
- self, cluster_name, region, zones, dryrun)
  return dict(
  cloud_specific_variables,
  **{

sky/templates/lambda-ray.yml.j2

@@ -5,6 +5,26 @@ max_workers: {{num_nodes - 1}}
 upscaling_speed: {{num_nodes - 1}}
 idle_timeout_minutes: 60
 
+{%- if docker_image is not none %}
+docker:
+ image: {{docker_image}}
+ container_name: {{docker_container_name}}
+ run_options:
+ - --ulimit nofile=1048576:1048576
+ {%- for run_option in docker_run_options %}
+ - {{run_option}}
+ {%- endfor %}
+ {%- if docker_login_config is not none %}
+ docker_login_config:
+ username: |-
+ {{docker_login_config.username}}
+ password: |-
+ {{docker_login_config.password}}
+ server: |-
+ {{docker_login_config.server}}
+ {%- endif %}
+{%- endif %}
+
 provider:
  type: external
  module: sky.provision.lambda

sky/utils/command_runner.py

@@ -502,8 +502,10 @@ def close_cached_connection(self) -> None:
  if self.ssh_control_name is not None:
  control_path = _ssh_control_path(self.ssh_control_name)
  if control_path is not None:
+ # Suppress the `Exit request sent.` output for this comamnd
+ # which would interrupt the CLI spinner.
  cmd = (f'ssh -O exit -S {control_path}/%C '
- f'{self.ssh_user}@{self.ip}')
+ f'{self.ssh_user}@{self.ip} > /dev/null 2>&1')
  logger.debug(f'Closing cached connection {control_path!r} with '
  f'cmd: {cmd}')
  log_lib.run_with_log(cmd,