[Core] Fix A10 GPU on Azure (#3707)

* init * works. todo: only do this for A10 VMs * only install for A10 instances * merge into one template * Update sky/skylet/providers/azure/node_provider.py Co-authored-by: Zhanghao Wu <[email protected]> * add warning * apply suggestions from code review * Update sky/clouds/azure.py Co-authored-by: Zhanghao Wu <[email protected]> --------- Co-authored-by: Zhanghao Wu <[email protected]>
skypilot-org · Jul 5, 2024 · 994d35a · 994d35a
1 parent 05ce5e9
commit 994d35a
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 8 deletions.
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
@@ -2020,8 +2020,16 @@ def provision_with_retries(
  failover_history: List[Exception] = list()
 
  style = colorama.Style
+ fore = colorama.Fore
  # Retrying launchable resources.
  while True:
+ if (isinstance(to_provision.cloud, clouds.Azure) and
+ to_provision.accelerators is not None and
+ 'A10' in to_provision.accelerators):
+ logger.warning(f'{style.BRIGHT}{fore.YELLOW}Trying to launch '
+ 'an A10 cluster on Azure. This may take ~20 '
+ 'minutes due to driver installation.'
+ f'{style.RESET_ALL}')
  try:
  # Recheck cluster name as the 'except:' block below may
  # change the cloud assignment.

diff --git a/sky/clouds/azure.py b/sky/clouds/azure.py
@@ -7,7 +7,7 @@
 import subprocess
 import textwrap
 import typing
-from typing import Dict, Iterator, List, Optional, Tuple
+from typing import Any, Dict, Iterator, List, Optional, Tuple
 
 import colorama
 
@@ -269,13 +269,12 @@ def get_vcpus_mem_from_instance_type(
  def get_zone_shell_cmd(cls) -> Optional[str]:
  return None
 
- def make_deploy_resources_variables(
- self,
- resources: 'resources.Resources',
- cluster_name_on_cloud: str,
- region: 'clouds.Region',
- zones: Optional[List['clouds.Zone']],
- dryrun: bool = False) -> Dict[str, Optional[str]]:
+ def make_deploy_resources_variables(self,
+ resources: 'resources.Resources',
+ cluster_name_on_cloud: str,
+ region: 'clouds.Region',
+ zones: Optional[List['clouds.Zone']],
+ dryrun: bool = False) -> Dict[str, Any]:
  assert zones is None, ('Azure does not support zones', zones)
 
  region_name = region.name
@@ -315,6 +314,10 @@ def make_deploy_resources_variables(
  'image_version': version,
  }
 
+ # Setup the A10 nvidia driver.
+ need_nvidia_driver_extension = (acc_dict is not None and
+ 'A10' in acc_dict)
+
  # Setup commands to eliminate the banner and restart sshd.
  # This script will modify /etc/ssh/sshd_config and add a bash script
  # into .bashrc. The bash script will restart sshd if it has not been
@@ -367,6 +370,7 @@ def _failover_disk_tier() -> Optional[resources_utils.DiskTier]:
  # Azure does not support specific zones.
  'zones': None,
  **image_config,
+ 'need_nvidia_driver_extension': need_nvidia_driver_extension,
  'disk_tier': Azure._get_disk_type(_failover_disk_tier()),
  'cloud_init_setup_commands': cloud_init_setup_commands,
  'azure_subscription_id': self.get_project_id(dryrun),

diff --git a/sky/skylet/providers/azure/node_provider.py b/sky/skylet/providers/azure/node_provider.py
@@ -303,6 +303,33 @@ def _create_node(self, node_config, tags, count):
  template_params["nsg"] = self.provider_config["nsg"]
  template_params["subnet"] = self.provider_config["subnet"]
 
+ if node_config.get("need_nvidia_driver_extension", False):
+ # Configure driver extension for A10 GPUs. A10 GPUs requires a
+ # special type of drivers which is available at Microsoft HPC
+ # extension. Reference: https://forums.developer.nvidia.com/t/ubuntu-22-04-installation-driver-error-nvidia-a10/285195/2
+ for r in template["resources"]:
+ if r["type"] == "Microsoft.Compute/virtualMachines":
+ # Add a nested extension resource for A10 GPUs
+ r["resources"] = [
+ {
+ "type": "extensions",
+ "apiVersion": "2015-06-15",
+ "location": "[variables('location')]",
+ "dependsOn": [
+ "[concat('Microsoft.Compute/virtualMachines/', parameters('vmName'), copyIndex())]"
+ ],
+ "name": "NvidiaGpuDriverLinux",
+ "properties": {
+ "publisher": "Microsoft.HpcCompute",
+ "type": "NvidiaGpuDriverLinux",
+ "typeHandlerVersion": "1.9",
+ "autoUpgradeMinorVersion": True,
+ "settings": {},
+ },
+ },
+ ]
+ break
+
  parameters = {
  "properties": {
  "mode": DeploymentMode.incremental,

diff --git a/sky/templates/azure-ray.yml.j2 b/sky/templates/azure-ray.yml.j2
@@ -80,6 +80,7 @@ available_node_types:
  # billingProfile:
  # maxPrice: -1
  {%- endif %}
+ need_nvidia_driver_extension: {{need_nvidia_driver_extension}}
  # TODO: attach disk
 {% if num_nodes > 1 %}
  ray.worker.default:
@@ -108,6 +109,7 @@ available_node_types:
  # billingProfile:
  # maxPrice: -1
  {%- endif %}
+ need_nvidia_driver_extension: {{need_nvidia_driver_extension}}
 {%- endif %}
 
 head_node_type: ray.head.default