[robustness] cover some potential resource leakage cases

[cg505]

Includes the following changes:

if a newly-created cluster is missing from the cloud, wait before deleting
This protects against leaking instances that haven't appeared yet, because they were just created.
Addresses #4431, see that issue for more details.

• This is observed to affect AWS if the timing is exactly right - there's a <1s window where the instances have not yet appeared.
• Haven't been able to reproduce this on any other clouds.

confirm cluster actually terminates before deleting from the db
Previously, we optimistically assumed that a successful termination request will always delete cloud instances. This change validates that the instances are actually terminating or terminated before deleting them from our state database. This protects against silent failures in the termination, which are not expected, but are hypothesized to be possible.

• This potential leak has not been reproduced.
• This does not add any additional delay to AWS or GCP. Azure takes around 10s to transition the instances, so cluster teardown now takes that much longer.

avoid deleting cluster data outside the primary provision loop
The main provision/failover/retry loop is quite complicated. There were some areas of code which could throw an exception and delete the cluster from our state database without terminating the instances. Attempt to clean up these paths.
The previous change (confirm cluster actually terminates before deleting from the db) also makes the provision loop safer. Since the instance termination and the state cleanup are in different places, this protects against the case where we don't successfully terminate the instances but still clean up our local state.

part of #4410

Tested (run the relevant ones):

• Code formatting: bash format.sh
• Any manual or new tests for this PR (please specify below)
  • Manually tried to reproduce leaks on AWS, Azure, and GCP. For successfully reproduced leaks (Race between status update and instance creation can cause resource leak #4431), verified the fix.
  • Normal launch and failover on these clouds.
  • sky launch --gpus H100:8 for longer failover checking
• Relevant individual smoke tests: pytest tests/test_smoke.py::test_minimal (on aws, azure, gcp)

[Michaelvll]

Thanks for the fix @cg505! A main comment I have is the place we do the wait for terminating / stopping state. I am leaning towards to have the cloud's specific instance.py to wait for the state transition, see:

skypilot/sky/provision/aws/instance.py

Lines 635 to 640 in 6e50832

	# TODO(suquark): Currently, the implementation of GCP and Azure will
	# wait util the cluster is fully terminated, while other clouds just
	# trigger the termination process (via http call) and then return.
	# It's not clear that which behavior should be expected. We will not
	# wait for the termination for now, since this is the default behavior
	# of most cloud implementations (including AWS).

vs

skypilot/sky/provision/gcp/instance.py

Lines 491 to 511 in 6e50832

	# Check if the instance is actually stopped.
	# GCP does not fully stop an instance even after
	# the stop operation is finished.
	for _ in range(constants.MAX_POLLS_STOP):
	handler_to_instances = _filter_instances(
	handler_to_instances.keys(),
	project_id,
	zone,
	label_filters,
	lambda handler: handler.NON_STOPPED_STATES,
	included_instances=all_instances,
	)
	if not handler_to_instances:
	break
	time.sleep(constants.POLL_INTERVAL)
	else:
	raise RuntimeError(f'Maximum number of polls: '
	f'{constants.MAX_POLLS_STOP} reached. '
	f'Instance {all_instances} is still not in '
	'STOPPED status.')

vs

skypilot/sky/provision/gcp/instance.py

Lines 567 to 570 in 6e50832

	# We don't wait for the instances to be terminated, as it can take a long
	# time (same as what we did in ray's node_provider).

[cg505]

A main comment I have is the place we do the wait for terminating / stopping state. I am leaning towards to have the cloud's specific instance.py to wait for the state transition

@Michaelvll A couple points:

1. I guess the fact that this is inconsistent between clouds is more reason to make the implementation cloud-agnostic. Then we don't need to go through each cloud implementation to make sure it's doing "the right thing". (Even though this would probably be good to have.) Fixing this requires a good understanding of the exact circumstances that allow for instance state transitions and outside of the big 3 clouds this is often poorly documented or not documented at all.
2. I think there is some confusion around waiting for the state change. Most clouds have an explicit stopping/terminating state, before reaching the ultimate stopped/terminating state. It's not clear if the comments you linked are talking about the former or the latter. For our use case, we basically think of stopping/terminating as equivalent to stopped/terminated, as discussed in the other thread.
   For instance, looking at the GCP code links:
   • For termination, the comment says "We don't wait for the instances to be terminated" but in my testing, the instances immediately transition to STOPPING, which we consider as "good enough" for our checks. (Since it's clear that our termination request succeeded at least.
   • For stopping, the comment says "GCP does not fully stop an instance even after the stop operation is finished." and the code will wait until one of STOPPED_STATES is reached (STOPPING_STATES is not good enough). But query_instances will consider STOPPING_STATES as ClusterStatus.STOPPED.

Ideally, we should definitely define, document, and enforce the exact invariants for various states and transitions, so that all the clouds have a clear indication of how it should work. But given we don't have that right now, I think this is the safest way.

P.S. If we're still concerned about the "wait for transition" change, we could make it only apply to termination, since it's much worse than stopping in terms of the leak.

[Michaelvll]

thanks @cg505! Once we have the two changes we mentioned offline, it should be good to go.

[Michaelvll]

Seems this throw and catch is not necessary? Should we just use a flag?

[cg505]

It won't re-throw if we should retry. Not sure what you mean by using a flag. never mind, I see

[Michaelvll]

Ooo, I mean exception handling is more expensive and here it seems unnecessary to raise as we can do something like the following:

while True:
  is_fail = False
  for node in nodes: 
    if condition:
        is_fail = True
        break
  if is_fail:
      ...

[Michaelvll]

Thanks @cg505! Left two comments for the comments in the code. Otherwise, LGTM.