fix: fix mig-partition timeout to ensure the service completes

[henryli001]

What type of PR is this?

What this PR does / why we need it:
mig-partition service is only given 30sec to restart which is not sufficient for the service to complete when the restart operation is executed the first time. This will cause the provisioning process to keep retrying (up to 100 times), which may eventually happen since the previous attempts to enable MIG mode have already MIG-enabled some GPUs and if not, a re-attempt of provisioning will happen. Eventually, even if the cluster gets successfully provisioned, the entire provisioning process ends up taking multiple hours in local experiments.

This behavior was captured on the 8-card A100 GPU SKUs such as https://learn.microsoft.com/en-us/azure/virtual-machines/sizes/gpu-accelerated/ndma100v4-series?tabs=sizebasic for both Azure Linux and Ubuntu AKS. Thru existing experiments, the first attempt of systemctl restart mig-partition would take around 2min+ on both Azure Linux and Ubuntu. Thus, this change increases the timeout value of the mig-partition service itself to 5min since the default timeout value of systemd is 90sec. The logic to perform systemctl restart is also modified to timeout in 5min for mig-partition service specifically while maintaing the original 30sec for other systemd services.

Which issue(s) this PR fixes:

Requirements:

• uses conventional commit messages
• includes documentation
• adds unit tests
• tested upgrade from previous version

Special notes for your reviewer:

Release note:

none

[cameronmeissner]

@henryli001 friendly reminder to please avoid including internal ADO links within PR descriptions/comments

[cameronmeissner]

I'd prefer bubbling the timeout parameter up to the top-level from systemctl_restart instead of needing to pivot based on service name

[henryli001]

refactored the implementation

[github-actions]

No changes to cached containers or packages on Windows VHDs

[ganeshkumarashok]

Justing noting that the PR description may not apply for Ubuntu nodes. I've not noticed MIG enablement in AKS Ubuntu nodes take "multiple hours in local experiments" when I've provisioned them. Curious why there's such a difference between OSes.

[ganeshkumarashok]

What's the default timeout when this is not added?

[henryli001]

The default timeout value for normal systemd service is 90s

[ganeshkumarashok]

Thanks for making this fix. Were you able to check whether this PR fixed the issue?

[henryli001]

Justing noting that the PR description may not apply for Ubuntu nodes. I've not noticed MIG enablement in AKS Ubuntu nodes take "multiple hours in local experiments" when I've provisioned them. Curious why there's such a difference between OSes.

The issue was also reproed on Ubuntu 2204 using VM size: Standard_ND96asr_v4

Dumping the failure message here:
[{"code":"VMExtensionProvisioningError","target":"0","message":"VM has reported a failure when processing extension 'vmssCSE' (publisher 'Microsoft.Azure.Extensions' and type 'CustomScript'). Error message: 'Enable failed: failed to execute command: command terminated with exit status=124\n[stdout]\n{ \"ExitCode\": \"124\", \"Output\": \"ontrol process exited, code=killed, status=15/TERM\\nMar 07 01:07:11 aks-nodepool1-12368757-vmss000000 systemd[1]: mig-partition.service: Failed with result 'signal'.\\nMar 07 01:07:11 aks-nodepool1-12368757-vmss000000 systemd[1]: Stopped Apply MIG configuration on Nvidia A100 GPU.\\nMar 07 01:07:11 aks-nodepool1-12368757-vmss000000 systemd[1]: mig-partition.service: Consumed 9.650s CPU time.\\n

[henryli001]

Thanks for making this fix. Were you able to check whether this PR fixed the issue?

Customized AgentBaker E2E using a Azure Linux AKS 3.0 image and was able to successfully deploy a MIG node and the node rebooted as expected and eventually became ready. Not able to test on Ubuntu since I don't have read permission on the sub that hosts the Ubuntu test image

[henryli001]

Thanks for making this fix. Were you able to check whether this PR fixed the issue?

Customized AgentBaker E2E using a Azure Linux AKS 3.0 image and was able to successfully deploy a MIG node and the node rebooted as expected and eventually became ready. Not able to test on Ubuntu since I don't have read permission on the sub that hosts the Ubuntu test image

Able to run test against Ubuntu using the pipeline and the MIG node got provisioned successfully. After the node got ready after a reboot which was expected on A100, nvidia-smi is showing the MIG devices correctly.

[github-actions]

No changes to cached containers or packages on Windows VHDs

[ganeshkumarashok]

Thanks Henry, this lgtm!