Fix AMD GPU driver installation failure on Ubuntu 22.04 (jammy)

src/VirtualClient/VirtualClient.Dependencies.UnitTests/AMDGPUDriverInstallationTests.cs

@@ -47,21 +47,57 @@ public void TearDown()
         }
 
         [Test]
-        public void AMDGPUDriverInstallationDependencyThrowsIfLinuxInstallationFileIsEmpty()
+        public void AMDGPUDriverInstallationDependencyThrowsIfNoInstallFileForUnsupportedCodename()
         {
-            this.SetupDefaultMockBehavior(PlatformID.Unix, string.Empty, string.Empty);
+            this.SetupDefaultMockBehavior(PlatformID.Unix, string.Empty, string.Empty, "unsupported_codename");
 
             DependencyException exc = Assert.ThrowsAsync<DependencyException>(() => this.component.ExecuteAsync(CancellationToken.None));
             Assert.AreEqual(ErrorReason.DependencyNotFound, exc.Reason);
         }
 
+        [Test]
+        public async Task AMDGPUDriverInstallationResolvesInstallFileFromMappingForJammy()
+        {
+            // No explicit LinuxInstallationFile — should resolve from built-in mapping for jammy
+            this.SetupDefaultMockBehavior(PlatformID.Unix, linuxInstallationFile: string.Empty, osVersionCodename: "jammy");
+
+            await this.component.ExecuteAsync(CancellationToken.None);
+            Assert.IsTrue(this.fixture.ProcessManager.CommandsExecuted(
+                "wget https://repo.radeon.com/amdgpu-install/6.3.3/ubuntu/jammy/amdgpu-install_6.3.60303-1_all.deb"));
+        }
+
+        [Test]
+        public async Task AMDGPUDriverInstallationResolvesInstallFileFromMappingForFocal()
+        {
+            // No explicit LinuxInstallationFile — should resolve from built-in mapping for focal
+            this.SetupDefaultMockBehavior(PlatformID.Unix, linuxInstallationFile: string.Empty, osVersionCodename: "focal");
+
+            await this.component.ExecuteAsync(CancellationToken.None);
+            Assert.IsTrue(this.fixture.ProcessManager.CommandsExecuted(
+                "wget https://repo.radeon.com/amdgpu-install/5.5/ubuntu/focal/amdgpu-install_5.5.50500-1_all.deb"));
+        }
+
+        [Test]
+        public async Task AMDGPUDriverInstallationUsesProfileInstallFileWhenProvided()
+        {
+            // Explicit LinuxInstallationFile provided — should use it regardless of codename
+            string customUrl = "https://repo.radeon.com/amdgpu-install/5.5/ubuntu/focal/amdgpu-install_5.5.50500-1_all.deb";
+            this.SetupDefaultMockBehavior(PlatformID.Unix, linuxInstallationFile: customUrl, osVersionCodename: "jammy");
+
+            await this.component.ExecuteAsync(CancellationToken.None);
+            Assert.IsTrue(this.fixture.ProcessManager.CommandsExecuted(
+                $"wget {customUrl}"));
+        }
+
         [Test]
         public async Task AMDGPUDriverInstallationDependencyStartsCorrectProcessesOnExecuteForLinux()
         {
             this.SetupDefaultMockBehavior(PlatformID.Unix);
 
             List<string> commands = new List<string>
             {
+                "sudo bash -c \"dpkg --remove --force-remove-reinstreq amdgpu-dkms || true\"",
+                "sudo bash -c \"dpkg --configure -a || true\"",
                 "apt-get -yq update",
                 "sudo apt-get install -yq libpci3 libpci-dev doxygen unzip cmake git",
                 "sudo apt-get install -yq libnuma-dev libncurses5",
@@ -111,7 +147,7 @@ public async Task AMDGPUDriverInstallationDependencyDoesNotInstallAMDGPUDriverIf
             Assert.IsFalse(this.fixture.ProcessManager.CommandsExecuted($"{installScriptPath} -INSTALL -OUTPUT screen"));
         }
 
-        private void SetupDefaultMockBehavior(PlatformID platformID, string gpuModel = "v620", string linuxInstallationFile = "https://repo.radeon.com/amdgpu-install/5.5/ubuntu/focal/amdgpu-install_5.5.50500-1_all.deb")
+        private void SetupDefaultMockBehavior(PlatformID platformID, string gpuModel = "v620", string linuxInstallationFile = "https://repo.radeon.com/amdgpu-install/5.5/ubuntu/focal/amdgpu-install_5.5.50500-1_all.deb", string osVersionCodename = "focal")
         {
             this.fixture.Setup(platformID);
             this.mockPackage = new DependencyPath("amddriverpackage", this.fixture.GetPackagePath("amddriverpackage"));
@@ -141,6 +177,17 @@ private void SetupDefaultMockBehavior(PlatformID platformID, string gpuModel = "
 
             this.fixture.File.Setup(file => file.Exists(It.IsAny<string>())).Returns(true);
 
+            // Mock /etc/os-release so the codename can be detected during InitializeAsync
+            if (platformID == PlatformID.Unix)
+            {
+                string osReleaseContent = string.IsNullOrEmpty(osVersionCodename)
+                    ? "NAME=\"Ubuntu\"\nVERSION=\"22.04 LTS\"\nID=ubuntu\nPRETTY_NAME=\"Ubuntu 22.04 LTS\""
+                    : $"NAME=\"Ubuntu\"\nVERSION_CODENAME={osVersionCodename}\nID=ubuntu\nPRETTY_NAME=\"Ubuntu LTS\"";
+
+                this.fixture.File.Setup(file => file.ReadAllTextAsync("/etc/os-release", It.IsAny<CancellationToken>()))
+                    .ReturnsAsync(osReleaseContent);
+            }
+
             this.fixture.SystemManagement.SetupGet(mgr => mgr.ProcessManager).Returns(this.mockProcessManager.Object);
 
             this.fixture.ApiClient.Setup(client => client.GetStateAsync(It.IsAny<string>(), It.IsAny<CancellationToken>(), It.IsAny<IAsyncPolicy<HttpResponseMessage>>()))

src/VirtualClient/VirtualClient.Dependencies/AMDGPUDriverInstallation.cs

@@ -9,6 +9,7 @@ namespace VirtualClient.Dependencies
     using System.IO;
     using System.IO.Abstractions;
     using System.Linq;
+    using System.Text.RegularExpressions;
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.CodeAnalysis;
@@ -26,11 +27,21 @@ public class AMDGPUDriverInstallation : VirtualClientComponent
     {
         private const string Mi25ExeName = "AMD-mi25.exe";
         private const string V620ExeName = "Setup.exe";
+
+        // Known-good ROCm installation URLs for each supported Ubuntu codename.
+        // These are the default URLs used when the profile does not specify a LinuxInstallationFile.
+        private static readonly Dictionary<string, string> SupportedInstallationFiles = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase)
+        {
+            { "focal", "https://repo.radeon.com/amdgpu-install/5.5/ubuntu/focal/amdgpu-install_5.5.50500-1_all.deb" },
+            { "jammy", "https://repo.radeon.com/amdgpu-install/6.3.3/ubuntu/jammy/amdgpu-install_6.3.60303-1_all.deb" }
+        };
+
         private IPackageManager packageManager;
         private IFileSystem fileSystem;
         private ISystemManagement systemManager;
         private IStateManager stateManager;
         private LinuxDistributionInfo linuxDistributionInfo;
+        private string osVersionCodename;
 
         /// <summary>
         /// Initializes a new instance of the <see cref="AMDGPUDriverInstallation"/> class.
@@ -185,12 +196,17 @@ protected override async Task InitializeAsync(EventContext telemetryContext, Can
             {
                 this.linuxDistributionInfo = await this.systemManager.GetLinuxDistributionAsync(cancellationToken)
                     .ConfigureAwait(false);
+
+                this.osVersionCodename = await this.DetectOsVersionCodenameAsync(cancellationToken)
+                    .ConfigureAwait(false);
             }
         }
 
         [SuppressMessage("StyleCop.CSharp.ReadabilityRules", "SA1118:Parameter should not span multiple lines", Justification = "Readability")]
         private async Task InstallAMDGPUDriverLinux(EventContext telemetryContext, CancellationToken cancellationToken)
         {
+            this.ResolveLinuxInstallationFile();
+
             if (string.IsNullOrWhiteSpace(this.LinuxInstallationFile))
             {
                 throw new DependencyException($"The linux installation file can not be null or empty and it is: {this.LinuxInstallationFile}", ErrorReason.DependencyNotFound);
@@ -199,7 +215,10 @@ private async Task InstallAMDGPUDriverLinux(EventContext telemetryContext, Cance
             // The .bashrc file is used to define commands that should be run whenever the system
             // is booted. For the purpose of the AMD GPU driver installation, we want to include extra
             // paths in the $PATH environment variable post installation.
-            string bashRcPath = $"/home/{this.Username}/.bashrc";
+            string userHome = this.GetUserHomePath();
+            string bashRcPath = $"{userHome}/.bashrc";
+
+            this.fileSystem.Directory.CreateDirectory(Path.GetDirectoryName(bashRcPath) !);
 
             this.fileSystem.Directory.CreateDirectory(Path.GetDirectoryName(bashRcPath) !);
 
@@ -255,6 +274,11 @@ private List<string> PrerequisiteCommands()
             switch (this.linuxDistributionInfo.LinuxDistribution)
             {
                 case LinuxDistribution.Ubuntu:
+                    // Clean up any broken package state from previous failed installation attempts.
+                    // The amdgpu-dkms package can be left in a half-configured state if the DKMS module
+                    // build fails, which causes all subsequent apt-get commands to fail.
+                    commands.Add("bash -c \"dpkg --remove --force-remove-reinstreq amdgpu-dkms || true\"");
+                    commands.Add("bash -c \"dpkg --configure -a || true\"");
                     commands.Add("apt-get -yq update");
                     commands.Add("apt-get install -yq libpci3 libpci-dev doxygen unzip cmake git");
                     commands.Add("apt-get install -yq libnuma-dev libncurses5");
@@ -293,12 +317,14 @@ private List<string> VersionSpecificInstallationCommands()
 
         private List<string> PostInstallationCommands()
         {
+            string userHome = this.GetUserHomePath();
+
             // last 2 command are to make sure that we are blacklisting AMD GPU drivers before rebooting
             return new List<string>
             {
                 "amdgpu-install -y --usecase=hiplibsdk,rocm,dkms",
                 $"bash -c \"echo 'export PATH=/opt/rocm/bin${{PATH:+:${{PATH}}}}' | " +
-                $"sudo tee -a /home/{this.Username}/.bashrc\"",
+                $"sudo tee -a {userHome}/.bashrc\"",
                 $"bash -c \"echo 'blacklist amdgpu' | sudo tee -a /etc/modprobe.d/amdgpu.conf \"",
                 "update-initramfs -u -k all"
             };
@@ -336,6 +362,67 @@ await this.LogProcessDetailsAsync(process, telemetryContext, "AMDGPUDriverInstal
             }
         }
 
+        /// <summary>
+        /// Reads /etc/os-release to detect the OS version codename (e.g., "focal", "jammy").
+        /// Follows the same pattern as MongoDBServerInstallation's codename detection.
+        /// </summary>
+        private async Task<string> DetectOsVersionCodenameAsync(CancellationToken cancellationToken)
+        {
+            try
+            {
+                string osReleaseContent = await this.fileSystem.File.ReadAllTextAsync("/etc/os-release", cancellationToken)
+                    .ConfigureAwait(false);
+
+                Match match = Regex.Match(osReleaseContent, @"VERSION_CODENAME=(\w+)", RegexOptions.Multiline);
+                return match.Success ? match.Groups[1].Value : null;
+            }
+            catch
+            {
+                // If /etc/os-release cannot be read, return null and let ResolveLinuxInstallationFile
+                // fall back to the profile-provided LinuxInstallationFile parameter.
+                return null;
+            }
+        }
+
+        /// <summary>
+        /// Resolves the correct Linux installation file URL based on the detected OS codename.
+        /// If the profile provides an explicit LinuxInstallationFile, that takes precedence.
+        /// Otherwise, the built-in mapping of codename to ROCm URL is used.
+        /// </summary>
+        private void ResolveLinuxInstallationFile()
+        {
+            // If the profile explicitly provided a LinuxInstallationFile, use it as-is.
+            if (!string.IsNullOrWhiteSpace(this.LinuxInstallationFile))
+            {
+                return;
+            }
+
+            // Look up the detected codename in the built-in mapping.
+            if (!string.IsNullOrWhiteSpace(this.osVersionCodename)
+                && AMDGPUDriverInstallation.SupportedInstallationFiles.TryGetValue(this.osVersionCodename, out string resolvedUrl))
+            {
+                this.LinuxInstallationFile = resolvedUrl;
+                return;
+            }
+
+            throw new DependencyException(
+                $"No AMD GPU driver installation file is available for the detected OS codename '{this.osVersionCodename}'. " +
+                $"Supported codenames: {string.Join(", ", AMDGPUDriverInstallation.SupportedInstallationFiles.Keys)}. " +
+                $"You can provide an explicit URL via the 'LinuxInstallationFile' profile parameter.",
+                ErrorReason.DependencyNotFound);
+        }
+
+        private string GetUserHomePath()
+        {
+            string username = this.Username;
+            if (string.Equals(username, "root", StringComparison.OrdinalIgnoreCase))
+            {
+                return "/root";
+            }
+
+            return $"/home/{username}";
+        }
+
         private async Task InstallAMDGPUDriverWindows(EventContext telemetryContext, CancellationToken cancellationToken)
         {
             string installerPath = string.Empty;