[Azure] Fix to sync NSG status while opening ports

[landscapepainter]

This resolves #3843.

Tested (run the relevant ones):

• Code formatting: bash format.sh
• Any manual or new tests for this PR (please specify below)
• All smoke tests: pytest tests/test_smoke.py --azure besides the ones failing from master branch.
• 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 @landscapepainter! Left two comments. : )

[Michaelvll]

Should we avoid assuming the nsg name for a cluster to be more robust to future changes, such #3845? We can just check all the NSG in the same resource group as we are using one resource group per cluster.

By doing that, let's keep the while and for loop in the same order as the original.

[landscapepainter]

@Michaelvll Thanks! I have two questions to address:

1. If I understand correctly, it doesn't seem like we can keep the original order of while and for loop. It seems necessary to call list_network_security_groups(resource_group) within the while loop (while not nsg_created:) so that the NSG status is updated on each iteration. This ensures that the check if nsg.provisioning_state not in ['Creating', 'Updating']: accurately reflects the latest state within the for loop (for nsg in nsg_list:). However, I might be overlooking something. Please let me know if that’s the case.

2. I can remove the check for specific nsg name from this PR, but this was also intended to be used for [Azure] Allow resource group specifiation for Azure instance provisioning #3764 as well. As you mentioned, we currently use one resource group per cluster, but this changes when resource group specification is used with sky serve. Serve controller and all the replicas would be created under the user specified resource group, and in this case, there would be multiple NSGs under the resource group. This was causing issues as opening ports while launching replica 4 was checking on NSG creation for replica 5 and got stuck.

   If losing track of the NSG name for the cluster throughout the launch cycle is your concern, we can keep on track of the NSG name while passing the name as a parameter to our ARM template instead of directly creating it from the ARM("nsgName": "[concat('sky-', parameters('clusterId'), '-nsg')]") and perhaps have that stored in provider_config as well, so we can read from it instead of using _NSG_NAME.format().

   2-1. Is it possible that we may need multiple NSGs for a single cluster in near future? We don't seem to do this now.

[Michaelvll]

For 1, can we check if there is an Azure API that get a single nsg instead of listing all of them, so that we can keep getting that nsg within the while loop? The current loop order is unintuitive, looping through the nsg list within a retry loop. It would be more readable to do the retry loop for each NSG (although we only select a single NSG).

For 2, it makes sense when a resource group can be specified. In that case, how about we just create a function in azure/config.py that generates the NSG name, so it can be reused in both config.py and here? We have to be careful with backward compatibility then, if we do the exact match of the NSG. How about we allow a list of NSG names to support legacy names, such as ray-{clusterId}-... and the one that will be added with #3848

[landscapepainter]

@Michaelvll Following your review, I found an API method .get to directly retrieve the nsg object. But there's an issue with using this with multiple possible NSG names(legacy names). Attempting to get an nsg instance with the following snippet takes about ~15s to return a ResourceNotFoundError when given a non-existent nsg name:

nsg = network_client.network_security_groups.get(resource_group_name, nsg_name)

And we have to provide non-existent nsg name at most twice with two legacy nsg names, which adds ~30s, to provision VM. So, I instead kept the listing method to iterate through the NSG list to find a matching name among our legacy and current nsg names. This is done in a separate method, _get_cluster_nsg, so it also resolves your concern on readability.

Also, following your comment, I created a method to create NSG name with our current naming convention and used it within config.py::bootstrap_instances and instance.py::_get_cluster_nsg, which is called by instance.py::open_ports.

These are done at befcbc1

[Michaelvll]

Thanks for the update @landscapepainter! Mostly looks good to me with some minor fixes for the function structure

[Michaelvll]

nit: We should prioritize the latest generation NSG name in the resource group.

[landscapepainter]

@Michaelvll I checked if the list method we are using supports a way to list the nsg objects following the creation time, but it does not support a way. You can search for 'def list' from this page if you want to take a look at the method: https://raw.githubusercontent.com/Azure/azure-sdk-for-python/7aeef6f433e98571b0b7624f57d3d89570b94949/sdk/network/azure-mgmt-network/azure/mgmt/network/v2024_01_01/aio/operations/_operations.py

Another approach we can take is to get the list and sort it from our end, but the nsg objects that are returned by list method does not contain any attributes related to creation time.

[landscapepainter]

Thanks @Michaelvll! This is ready for another look.