[k8s] Multi-node support for Kubernetes

[romilbhardwaj]

Adds Kubernetes multi-node support. It was fairly straightforward, since most of the implementation was already done in #2096. Main thing to note is that spreading pods across nodes is set as a soft-constraint (i.e., try to place multi-node pods on different physical nodes. If not possible, ok to put them on the same node.)

Also closes #2628. This is a tricky one, since the problem is still not clear why it happens. See the issue for more details. After spending time investigating it, I have added a randomized sleep after kubectl port-forward is run.

While this is not the solution I wanted, it appears to fix the problem reliably. Ideally I would have wanted to get a signal from kubectl port-forward to know when it is ready, but it emits no signals to indicate it is actually ready (we already wait for it to print the expected output and use nc to verify the port is open).

What's working for multi-node:

• Basic multi-node launching, tasks and ssh (sky launch --num-nodes 2 -- echo hi)
• IP address fetching/populating SKYPILOT_NODE_IPS correctly
• Multi-node GPU jobs (pytorch distributed, tf examples)
  • Tested Nemo multi-node training for ~12 hours from [Examples] NeMo distributed training for BERT and GPT3 #2533

Tested (run the relevant ones):

• Code formatting: bash format.sh
• Tested SSH concurrency with the script from [k8s] SSH ProxyCommand script is not concurrency-safe #2628 and running run_command_in_parallel(['ssh','myclus','echo hi'], 20, 20) 10 times.
• Multi-node smoke tests are uncommented now - ran pytest tests/test_smoke.py --kubernetes -k "not TestStorageWithCredentials"

[romilbhardwaj]

This is now ready for a look! Also includes a fix for #2628 - updated the PR description above.

[Michaelvll]

This is awesome @romilbhardwaj! It is a surprisingly small amount of changes to make it work with multiple nodes. Tested with a GKE cluster and it works smoothly. One minor UX issue:

1. sky launch --gpus V100:0.5 --num-nodes 3 "hostname -i"
2. sky status: the resources column shows 3x Kubernetes(2CPU--8GB--1V100, {'V100': 1}), instead of V100:0.5, which is a bit surprising.

[Michaelvll]

Is a random wait time enough or should this be a socat and retry with a random backoff?

[Michaelvll]

Also it would be nice to have a newline at the EOF.

[romilbhardwaj]

Unfortunately, the socat command succeeds and starts running without any error even though the underlying kubectl port-forward may not be functional. Since there's no signal that I can use to detect if kubectl port-forward is ready, I used random wait which has been sufficient in my testing.

[Michaelvll]

Ahh, makes sense. I think it should be fine to wait for a random gap. Just curious, if something that checks the socket connection would work for checking if the connection gets setup correctly, like the following:

skypilot/sky/provision/provisioner.py

Lines 239 to 258 in 029f886

	def _wait_ssh_connection_direct(
	ip: str,
	ssh_user: str,
	ssh_private_key: str,
	ssh_control_name: Optional[str] = None,
	ssh_proxy_command: Optional[str] = None) -> bool:
	del ssh_control_name
	assert ssh_proxy_command is None, 'SSH proxy command is not supported.'
	try:
	with socket.create_connection((ip, 22), timeout=1) as s:
	if s.recv(100).startswith(b'SSH'):
	return True
	except socket.timeout: # this is the most expected exception
	pass
	except Exception: # pylint: disable=broad-except
	pass
	command = _ssh_probe_command(ip, ssh_user, ssh_private_key,
	ssh_proxy_command)
	logger.debug(f'Waiting for SSH to {ip}. Try: '
	f'{_shlex_join(command)}')

[Michaelvll]

nvm, it seems we have already done the nc check before this wait. Waiting looks good to me.

[romilbhardwaj]

Thanks for the review @Michaelvll! I've filed the fractional GPU bug in #2655, we should get to that too soon..