Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Core] Fix A10 GPU on Azure #3707

Merged
merged 9 commits into from
Jul 5, 2024
+49 −8
Merged

[Core] Fix A10 GPU on Azure #3707

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
3351caa
init
cblmemo Jun 28, 2024
74b669c
works. todo: only do this for A10 VMs
cblmemo Jun 29, 2024
d4a2f41
only install for A10 instances
cblmemo Jun 29, 2024
823ed8d
merge into one template
cblmemo Jul 2, 2024
64126d4
Merge remote-tracking branch 'origin/master' into fix-az-a10
cblmemo Jul 2, 2024
4a4789a
Update sky/skylet/providers/azure/node_provider.py
cblmemo Jul 2, 2024
4ccedce
add warning
cblmemo Jul 3, 2024
018181f
apply suggestions from code review
cblmemo Jul 4, 2024
f7b9935
Update sky/clouds/azure.py
cblmemo Jul 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions sky/backends/cloud_vm_ray_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
32 changes: 32 additions & 0 deletions sky/skylet/providers/azure/node_provider.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -303,6 +303,38 @@ def _create_node(self, node_config, tags, count):
template_params["nsg"] = self.provider_config["nsg"]
template_params["subnet"] = self.provider_config["subnet"]

# pylint: disable=import-outside-toplevel
from sky.clouds.service_catalog import azure_catalog

instance_type = node_config["azure_arm_parameters"].get("vmSize", "")
accs = azure_catalog.get_accelerators_from_instance_type(instance_type)
if accs is not None and "A10" in accs:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done! PTAL 🫡

# 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": {},
},
},
]
cblmemo marked this conversation as resolved.
Show resolved Hide resolved
break

parameters = {
"properties": {
"mode": DeploymentMode.incremental,
Expand Down
Loading