[Core] Fix A10 GPU on Azure

[cblmemo]

Fixes #3651. A10 GPUs require a special type of driver, Grid Driver, which is suggested to be installed by enabling the NvidiaGpuDriverLinux extension. This PR added this extension for A10 instances.

Tested (run the relevant ones):

• Code formatting: bash format.sh
• Any manual or new tests for this PR (please specify below)
  • The YAML in the issue and nvidia-smi works well.
  • The YAML in the issue, Qwen model e2e.
  • sky launch --cloud azure -c az-wo-a10 and checked that the drivers is not installed
• 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

[cblmemo]

To discuss: Should we put the template into a separate file like https://github.com/skypilot-org/skypilot/blob/master/sky/skylet/providers/azure/azure-config-template.json

[Michaelvll]

We should merge the parameters for A10 to original parameters to avoid two create_or_update which can cause significant overhead? Please help profile it.

[cblmemo]

Good point! I benchmarked and it reduces the provision time from 15m15s to 14m36s.

[cblmemo]

Since now we got more quotas, I also tested the YAML in the issue and it works in e2e as well 🫡

[Michaelvll]

Thanks for the update @cblmemo!

[Michaelvll]

Is there a way to avoid using azure_catalog in the node_provider as that seems to be an abstraction leakage, but I am ok to keep it as is for now.

[Michaelvll]

Actually, how about in the azure-ray.yaml we pass a need_nvidia_driver_extension: true within clouds/azure.py::make_deployable_variables. There we can check the accelerators, and get rid of the use of catalog here.

[cblmemo]

Done! PTAL 🫡

[Michaelvll]

Thanks for the update @cblmemo! Please check the comment above. We can get rid of the use of catalog in the node_provider with that. Other parts look good to me.