[Provisioner] Support docker in Lambda Cloud and TPU

[cblmemo]

Support docker as runtime on Lambda Cloud and TPU. TPU part is inherit from #3819.

Tested (run the relevant ones):

• Code formatting: bash format.sh
• Any manual or new tests for this PR (please specify below)
  • sky launch --image-id docker:continuumio/miniconda3:4.12.0 --cloud lambda
  • sky launch --image-id docker:continuumio/miniconda3:4.12.0 --cloud azure
• All smoke tests: pytest tests/test_smoke.py
• Relevant individual smoke tests: pytest tests/test_smoke.py::test_fill_in_the_name
• Backward compatibility tests: conda deactivate; bash -i tests/backward_compatibility_tests.sh

[Michaelvll]

Thanks for adding the support @cblmemo! Just want to make sure that id -nG $USER | grep -qw docker can be successfully run on other clouds as well. Can we test with AWS CPU/GPU, GCP CPU/GPU/TPU to make sure it does not break existing services?

[Michaelvll]

Can we keep this command running within other commands, e.g., _check_docker_installed, to reduce the ssh connection?

[Michaelvll]

Also, just curious, how does it compare with using sudo docker in all places docker is called. Asking because had the same issue before running docker on TPU VMs.

[cblmemo]

Lets merge this version first to unblock our user ;) filed an issue #4121 to keep track of this.

[cblmemo]

Thanks for adding the support @cblmemo! Just want to make sure that id -nG $USER | grep -qw docker can be successfully run on other clouds as well. Can we test with AWS CPU/GPU, GCP CPU/GPU/TPU to make sure it does not break existing services?

Do you mean testing them with docker? As the codepath changed by this PR won't be touched by non-docker clusters. I'm a little bit confused as we does not support docker with TPU for now?

[Michaelvll]

Thanks for adding the support @cblmemo! Just want to make sure that id -nG $USER | grep -qw docker can be successfully run on other clouds as well. Can we test with AWS CPU/GPU, GCP CPU/GPU/TPU to make sure it does not break existing services?

Do you mean testing them with docker? As the codepath changed by this PR won't be touched by non-docker clusters. I'm a little bit confused as we does not support docker with TPU for now?

Yes, testing with docker as runtime. Right, docker on TPU was not supported, but IIRC, it was mainly blocked by the exact issue you fixed, see here: #3819

Maybe we can pick that PR up once this is merged.

[cblmemo]

@Michaelvll Just tested all and it works well. On Lambda it failed to recognize the nvidia runtime for some reason, so i manually patched the run options. PTAL again!

skypilot/sky/provision/docker_utils.py

Lines 375 to 393 in 71a95f4

	def _configure_runtime(self, run_options: List[str]) -> List[str]:
	if self.docker_config.get('disable_automatic_runtime_detection'):
	return run_options
	runtime_output = (self._run(f'{self.docker_cmd} ' +
	'info -f "{{.Runtimes}}"'))
	if 'nvidia-container-runtime' in runtime_output:
	try:
	self._run('nvidia-smi', log_err_when_fail=False)
	return run_options + ['--runtime=nvidia', '--gpus all']
	except Exception as e: # pylint: disable=broad-except
	logger.debug(
	'Nvidia Container Runtime is present in the docker image'
	'specified, but no GPUs found on the cluster. It should '
	'still if the cluster is expected to have no GPU.\n'
	f' Details for nvidia-smi: {e}')
	return run_options
	return run_options

[Michaelvll]

Thanks for the update and adding TPU docker support @cblmemo! LGTM.